content: {-|
    Copyright  :  (C) 2012-2016, University of Twente,
                      2017     , Myrtle Software Ltd,
                      2017-2018, Google Inc.
                      2020-2022, QBayLogic B.V.
                      2022     , Google Inc.
    License    :  BSD2 (see the file LICENSE)
    Maintainer :  QBayLogic B.V. <devops@qbaylogic.com>
  
    Functions for expression manipulation
  -}
  
  {-# LANGUAGE MagicHash #-}
  {-# LANGUAGE TemplateHaskell #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE OverloadedStrings #-}
  
  module Clash.Netlist.Expr where
  
  import Control.Monad (zipWithM)
  import Control.Exception (assert)
  import Data.Set (fromList, member)
  import Data.Bits (Bits, testBit, setBit, zeroBits)
  import Data.Foldable (fold)
  import Data.Tree (Tree(..))
  import GHC.Stack (HasCallStack)
  import Data.Text (unpack)
  import Language.Haskell.TH.Quote (dataToPatQ)
  
  import qualified Clash.Sized.Vector as V (replicate)
  import qualified Clash.Sized.Internal.Index as I (fromInteger#)
  import qualified Clash.Sized.Internal.Signed as S (fromInteger#)
  import qualified Clash.Sized.Internal.Unsigned as U (fromInteger#)
  import qualified Clash.Sized.Internal.BitVector as BV
    (high, low, fromInteger#, fromInteger##)
  import GHC.Int (Int8(I8#), Int16(I16#), Int32(I32#), Int64(I64#))
  import GHC.Word (Word8(W8#), Word16(W16#), Word32(W32#), Word64(W64#))
  
  import Clash.Primitives.DSL (tySize)
  import Clash.Netlist.Types
    ( Size, Bit(..), Expr(..), HWType(..), Literal(..), Modifier(..)
    , BlackBoxContext(..)
    )
  
  -- | Turns a constant expression of known bitsize into their
  -- corresponding bitstream representation, arranged as a tree
  -- that corresponds to the structure of the expression.
  --
  -- NOTE: This conversion serves as a best effort approach and can be
  -- considered a hack. Fully featured constant expression evaluation is
  -- not available in clash yet and will replace this implementation
  -- once it is officially supported.
  bits :: HasCallStack => Size -> Expr -> Either Expr (Tree [Bool])
  bits size expr = case expr of
    Literal _ lit -> case lit of
      BitLit bLit   -> case bLit of
        H -> leaf [True]
        L -> leaf [False]
        _ -> Left expr
      BoolLit bLit      -> leaf [bLit]
      NumLit nLit       -> leaf $ toBits size nLit
      BitVecLit _ bvLit -> leaf $ toBits size bvLit
      VecLit lits       ->
        mapM (bits (size `div` length lits) . lit2Expr) lits >>= inner
      StringLit{}       -> Left expr
    DataCon ty m subs -> assert (tySize ty == size) $ case ty of
      Vector s t      -> vecBits (tySize t) s subs
      Product _ _ tys -> zipWithM bits (map tySize tys) subs >>= inner
      Sum _ cs        -> spBits expr size m subs $ map (const []) cs
      SP _ xs         -> spBits expr size m subs $ map (map tySize . snd) xs
      _               -> case subs of
        [e] -> bits size e
        []  -> leaf []
        _   -> Left expr
    -- appears in case of complex transformations, e.g.,
    -- >>> (bv2v 0b010) :: Vec 3 Bit
    -- >>> (iterate (SNat @3) not True) :: Vec 3 Bool
    -- >>> (complement <$> (True :> False :> Nil)) :: Vec 2 Bool
    Identifier{} -> Left expr
    DataTag{} -> Left expr
    BlackBoxE bbName _ _ _ _ bbCtx _ -> case unpack bbName of
      $(dataToPatQ (const Nothing) $ show 'BV.low) -> leaf [False]
      $(dataToPatQ (const Nothing) $ show 'BV.high) -> leaf [True]
      $(dataToPatQ (const Nothing) $ show 'V.replicate) -> case bbInputs bbCtx of
        [ (eSize, ty, _), (eValue, _, _) ] -> do
          bs <- bits (tySize ty) eSize
          let s = fromBits $ fold bs
          v <- bits (size `div` s) eValue
          inner $ replicate s v
        _ -> Left expr
      _ ->
        if unpack bbName `member` skippableBBs
        then skippableBBBits expr bbCtx size
        else Left expr
    ToBv _ _ e -> bits size e
    FromBv _ _ e -> bits size e
    IfThenElse cond match alt -> case bits 1 cond of
      Right (Node [True] [])  -> bits size match
      Right (Node [False] []) -> bits size alt
      _ -> Left expr
    Noop -> leaf []
  
   where
    -- known skippable blackboxes
    skippableBBs = fromList $ map show
      [ 'I.fromInteger#, 'S.fromInteger#, 'U.fromInteger#
      , 'BV.fromInteger#, 'BV.fromInteger##
      , 'I8#, 'I16#, 'I32#, 'I64#
      , 'W8#, 'W16#, 'W32#, 'W64#
      ]
  
    -- skips the blackbox conversion and obtains the constant result
    -- directly from the last input argument instead
    skippableBBBits e Context{..} n = case reverse bbInputs of
      (x, _, _) : _ -> bits n x
      _             -> Left e
  
    -- turns sum (& product) expressions into bitstreams (preserving the
    -- expressions' tree layout)
    spBits :: Expr -> Size -> Modifier -> [Expr] -> [[Size]]
           -> Either Expr (Tree [Bool])
    spBits e n m es sizes = case m of
      DC (_, i) -> do
        xs <- zipWithM bits (sizes !! i) es
        bs <- fold <$> inner xs
        l <- leaf $ toBits (n - length bs) i
        r <- leaf bs
        inner [ l, r ]
      _ -> Left e
  
    -- turns vector expressions into bitstream (preserving the
    -- expressions' tree layout)
    vecBits :: Size -> Int -> [Expr] -> Either Expr (Tree [Bool])
    vecBits elemSize elems = \case
      []   -> assert (elems == 0) $ leaf []
      x:xr -> assert (elems > 0) $ do
        (processedElems, cur) <- case x of
          DataCon t _ xs -> case t of
            Vector subElems (tySize -> subTySize) ->
              assert (subElems <= elems && subTySize == elemSize)
                ((subElems,) <$> vecBits elemSize subElems xs)
            _ -> (1,) <$> bits elemSize x
          _ -> (1,) <$> bits elemSize x
        sub <- vecBits elemSize (elems - processedElems) xr
        inner [cur, sub]
  
    -- creates a leaf node holding the leaf value
    leaf :: [a] -> Either b (Tree [a])
    leaf x = return $ Node x []
  
    -- creates an inner node (holding no value) with the given
    -- sub-trees
    inner :: [Tree [a]] -> Either b (Tree [a])
    inner = return . Node []
  
    -- turns a literal into an expression
    lit2Expr = Literal Nothing
  
  -- | Turns values into bitstreams of known length. If the bit stream
  -- requires more bits for representing the given value, then only the
  -- suffix of the corresponding bitstream gets returned.
  toBits :: Bits a => Int -> a -> [Bool]
  toBits n x =
    map (testBit x) [n-1,n-2..0]
  
  -- | Turns bitstreams into values.
  fromBits :: Bits a => [Bool] -> a
  fromBits xs =
    foldl setBit zeroBits $ map snd $ filter fst $ zip xs [n-1,n-2..0]
   where
    n = length xs
  