{-|
  Copyright  :  (C) 2015-2016, University of Twente,
                    2016     , Myrtle Software Ltd,
                    2021-2024, QBayLogic B.V.
  License    :  BSD2 (see the file LICENSE)
  Maintainer :  QBayLogic B.V. <devops@qbaylogic.com>

  Reductions of primitives

  Currently, it contains reductions for:

    * Clash.Sized.Vector.map
    * Clash.Sized.Vector.zipWith
    * Clash.Sized.Vector.traverse#
    * Clash.Sized.Vector.foldr
    * Clash.Sized.Vector.fold
    * Clash.Sized.Vector.dfold
    * Clash.Sized.Vector.(++)
    * Clash.Sized.Vector.head
    * Clash.Sized.Vector.tail
    * Clash.Sized.Vector.unconcatBitVector#
    * Clash.Sized.Vector.replicate
    * Clash.Sized.Vector.imap
    * Clash.Sized.Vector.dtfold
    * Clash.Sized.RTree.tfold
    * Clash.Sized.Vector.reverse
    * Clash.Sized.Vector.unconcat

  Partially handles:

    * Clash.Sized.Vector.transpose
-}

{-# LANGUAGE CPP #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE QuasiQuotes #-}

module Clash.Normalize.PrimitiveReductions where

import qualified Control.Lens                     as Lens
import           Control.Lens                     ((.=))
import           Control.Monad.Trans.Class        (lift)
import           Control.Monad.Trans.Maybe        (MaybeT (..))
import           Data.Bifunctor                   (second)
import           Data.List                        (mapAccumR)
import           Data.List.Extra                  (zipEqual)
#if MIN_VERSION_base(4,20,0)
import qualified Data.List.NonEmpty               as NE hiding (unzip)
import qualified Data.Functor                     as NE
#else
import qualified Data.List.NonEmpty               as NE
#endif
import qualified Data.Maybe                       as Maybe
import           Data.Semigroup                   (sconcat)
import           Data.Text.Extra                  (showt)
import           GHC.Stack                        (HasCallStack)

#if MIN_VERSION_ghc(9,0,0)
import           GHC.Builtin.Names
  (boolTyConKey, typeNatAddTyFamNameKey, typeNatMulTyFamNameKey,
   typeNatSubTyFamNameKey)
import           GHC.Types.SrcLoc                 (wiredInSrcSpan)
#else
import           PrelNames
  (boolTyConKey, typeNatAddTyFamNameKey, typeNatMulTyFamNameKey,
   typeNatSubTyFamNameKey)
import           SrcLoc                           (wiredInSrcSpan)
#endif

import           Clash.Core.DataCon               (DataCon)
import           Clash.Core.HasType
import           Clash.Core.Literal               (Literal (..))
import           Clash.Core.Name
  (nameOcc, Name(..), NameSort(User), mkUnsafeSystemName)
import           Clash.Core.Pretty                (showPpr)
import           Clash.Core.Term
  (IsMultiPrim (..), CoreContext (..), PrimInfo (..), Term (..), WorkInfo (..), Pat (..),
   collectTermIds, mkApps, PrimUnfolding(..))
import           Clash.Core.Type                  (LitTy (..), Type (..),
                                                   TypeView (..), coreView1,
                                                   mkFunTy, mkTyConApp,
                                                   splitFunForallTy, tyView)
import           Clash.Core.TyCon
  (TyConMap, TyConName, tyConDataCons, tyConName)
import           Clash.Core.TysPrim
  (integerPrimTy, typeNatKind, liftedTypeKind)
import           Clash.Core.Util
  (appendToVec, extractElems, extractTElems, mkRTree,
   mkUniqInternalId, mkUniqSystemTyVar, mkVec, dataConInstArgTys, primCo)
import           Clash.Core.Var                   (mkTyVar, mkLocalId)
import           Clash.Core.VarEnv                (extendInScopeSetList)
import qualified Clash.Data.UniqMap as UniqMap
import qualified Clash.Normalize.Primitives as NP (undefined)
import {-# SOURCE #-} Clash.Normalize.Strategy
import           Clash.Normalize.Types
import           Clash.Rewrite.Types
import           Clash.Rewrite.Util
import           Clash.Unique (fromGhcUnique)
import           Clash.Util
import qualified Clash.Util.Interpolate           as I

typeNatAdd :: TyConName
typeNatAdd :: TyConName
typeNatAdd =
  NameSort -> Text -> Unique -> SrcSpan -> TyConName
forall a. NameSort -> Text -> Unique -> SrcSpan -> Name a
Name NameSort
User Text
"GHC.TypeNats.+" (Unique -> Unique
fromGhcUnique Unique
typeNatAddTyFamNameKey) SrcSpan
wiredInSrcSpan

typeNatMul :: TyConName
typeNatMul :: TyConName
typeNatMul =
  NameSort -> Text -> Unique -> SrcSpan -> TyConName
forall a. NameSort -> Text -> Unique -> SrcSpan -> Name a
Name NameSort
User Text
"GHC.TypeNats.*" (Unique -> Unique
fromGhcUnique Unique
typeNatMulTyFamNameKey) SrcSpan
wiredInSrcSpan

typeNatSub :: TyConName
typeNatSub :: TyConName
typeNatSub =
  NameSort -> Text -> Unique -> SrcSpan -> TyConName
forall a. NameSort -> Text -> Unique -> SrcSpan -> Name a
Name NameSort
User Text
"GHC.TypeNats.-" (Unique -> Unique
fromGhcUnique Unique
typeNatSubTyFamNameKey) SrcSpan
wiredInSrcSpan

vecHeadPrim
  :: TyConName
  -- ^ Vec TyCon name
  -> Term
vecHeadPrim :: TyConName -> Term
vecHeadPrim TyConName
vecTcNm =
 -- head :: Vec (n+1) a -> a
  PrimInfo -> Term
Prim (Text
-> Type -> WorkInfo -> IsMultiPrim -> PrimUnfolding -> PrimInfo
PrimInfo Text
"Clash.Sized.Vector.head" (TyConName -> Type
vecHeadTy TyConName
vecTcNm) WorkInfo
WorkNever IsMultiPrim
SingleResult PrimUnfolding
NoUnfolding)

vecLastPrim
  :: TyConName
  -- ^ Vec TyCon name
  -> Term
vecLastPrim :: TyConName -> Term
vecLastPrim TyConName
vecTcNm =
  -- last :: Vec (n+1) a -> a
  -- has the same type signature as head, hence we're reusing its type
  -- definition here.
  PrimInfo -> Term
Prim (Text
-> Type -> WorkInfo -> IsMultiPrim -> PrimUnfolding -> PrimInfo
PrimInfo Text
"Clash.Sized.Vector.last" (TyConName -> Type
vecHeadTy TyConName
vecTcNm) WorkInfo
WorkNever IsMultiPrim
SingleResult PrimUnfolding
NoUnfolding)

vecHeadTy
  :: TyConName
  -- ^ Vec TyCon name
  -> Type
vecHeadTy :: TyConName -> Type
vecHeadTy TyConName
vecNm =
  TyVar -> Type -> Type
ForAllTy TyVar
nTV (Type -> Type) -> Type -> Type
forall a b. (a -> b) -> a -> b
$
  TyVar -> Type -> Type
ForAllTy TyVar
aTV (Type -> Type) -> Type -> Type
forall a b. (a -> b) -> a -> b
$
  Type -> Type -> Type
mkFunTy
    (TyConName -> [Type] -> Type
mkTyConApp TyConName
vecNm [TyConName -> [Type] -> Type
mkTyConApp TyConName
typeNatAdd [TyVar -> Type
VarTy TyVar
nTV, LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
1)], TyVar -> Type
VarTy TyVar
aTV])
    (TyVar -> Type
VarTy TyVar
aTV)
 where
  aTV :: TyVar
aTV = Type -> TyName -> TyVar
mkTyVar Type
liftedTypeKind (Text -> Unique -> TyName
forall a. Text -> Unique -> Name a
mkUnsafeSystemName Text
"a" Unique
0)
  nTV :: TyVar
nTV = Type -> TyName -> TyVar
mkTyVar Type
typeNatKind (Text -> Unique -> TyName
forall a. Text -> Unique -> Name a
mkUnsafeSystemName Text
"n" Unique
1)

vecTailPrim
  :: TyConName
  -- ^ Vec TyCon name
  -> Term
vecTailPrim :: TyConName -> Term
vecTailPrim TyConName
vecTcNm =
  -- tail :: Vec (n + 1) a -> Vec n a
  PrimInfo -> Term
Prim (Text
-> Type -> WorkInfo -> IsMultiPrim -> PrimUnfolding -> PrimInfo
PrimInfo Text
"Clash.Sized.Vector.tail" (TyConName -> Type
vecTailTy TyConName
vecTcNm) WorkInfo
WorkNever IsMultiPrim
SingleResult PrimUnfolding
NoUnfolding)

vecInitPrim
  :: TyConName
  -- ^ Vec TyCon name
  -> Term
vecInitPrim :: TyConName -> Term
vecInitPrim TyConName
vecTcNm =
  -- init :: Vec (n + 1) a -> Vec n a
  -- has the same type signature as tail, hence we're reusing its type
  -- definition here.
  PrimInfo -> Term
Prim (Text
-> Type -> WorkInfo -> IsMultiPrim -> PrimUnfolding -> PrimInfo
PrimInfo Text
"Clash.Sized.Vector.init" (TyConName -> Type
vecTailTy TyConName
vecTcNm) WorkInfo
WorkNever IsMultiPrim
SingleResult PrimUnfolding
NoUnfolding)

vecTailTy
  :: TyConName
  -- ^ Vec TyCon name
  -> Type
vecTailTy :: TyConName -> Type
vecTailTy TyConName
vecNm =
  TyVar -> Type -> Type
ForAllTy TyVar
nTV (Type -> Type) -> Type -> Type
forall a b. (a -> b) -> a -> b
$
  TyVar -> Type -> Type
ForAllTy TyVar
aTV (Type -> Type) -> Type -> Type
forall a b. (a -> b) -> a -> b
$
  Type -> Type -> Type
mkFunTy
    (TyConName -> [Type] -> Type
mkTyConApp TyConName
vecNm [TyConName -> [Type] -> Type
mkTyConApp TyConName
typeNatAdd [TyVar -> Type
VarTy TyVar
nTV, LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
1)], TyVar -> Type
VarTy TyVar
aTV])
    (TyConName -> [Type] -> Type
mkTyConApp TyConName
vecNm [TyVar -> Type
VarTy TyVar
nTV, TyVar -> Type
VarTy TyVar
aTV])
 where
  nTV :: TyVar
nTV = Type -> TyName -> TyVar
mkTyVar Type
typeNatKind (Text -> Unique -> TyName
forall a. Text -> Unique -> Name a
mkUnsafeSystemName Text
"n" Unique
0)
  aTV :: TyVar
aTV = Type -> TyName -> TyVar
mkTyVar Type
liftedTypeKind (Text -> Unique -> TyName
forall a. Text -> Unique -> Name a
mkUnsafeSystemName Text
"a" Unique
1)

-- | Makes two case statements: the first one extract the _head_ from the given
-- vector, the latter the tail.
extractHeadTail
  :: DataCon
  -- ^ The Cons (:>) constructor
  -> Type
  -- ^ Element type
  -> Integer
  -- ^ Length of the vector, must be positive
  -> Term
  -- ^ Vector to extract head from
  -> (Term, Term)
  -- ^ (head of vector, tail of vector)
extractHeadTail :: DataCon -> Type -> Integer -> Term -> (Term, Term)
extractHeadTail DataCon
consCon Type
elTy Integer
n Term
vec =
  case DataCon -> [Type] -> Maybe [Type]
dataConInstArgTys DataCon
consCon [Type]
tys of
    Just [Type
coTy, Type
_elTy, Type
restTy] ->
      let
        mTV :: TyVar
mTV = Type -> TyName -> TyVar
mkTyVar Type
typeNatKind (Text -> Unique -> TyName
forall a. Text -> Unique -> Name a
mkUnsafeSystemName Text
"m" Unique
0)
        co :: Id
co = Type -> TmName -> Id
mkLocalId Type
coTy (Text -> Unique -> TmName
forall a. Text -> Unique -> Name a
mkUnsafeSystemName Text
"_co_" Unique
1)
        el :: Id
el = Type -> TmName -> Id
mkLocalId Type
elTy (Text -> Unique -> TmName
forall a. Text -> Unique -> Name a
mkUnsafeSystemName Text
"el" Unique
2)
        rest :: Id
rest = Type -> TmName -> Id
mkLocalId Type
restTy (Text -> Unique -> TmName
forall a. Text -> Unique -> Name a
mkUnsafeSystemName Text
"res" Unique
3)

        pat :: Pat
pat = DataCon -> [TyVar] -> [Id] -> Pat
DataPat DataCon
consCon [TyVar
mTV] [Id
co, Id
el, Id
rest]
      in
        ( Term -> Type -> [Alt] -> Term
Case Term
vec Type
elTy [(Pat
pat, Id -> Term
Var Id
el)]
        , Term -> Type -> [Alt] -> Term
Case Term
vec Type
restTy [(Pat
pat, Id -> Term
Var Id
rest)] )
    Maybe [Type]
_ -> [Char] -> (Term, Term)
forall a. HasCallStack => [Char] -> a
error [Char]
"extractHeadTail: failed to instantiate Cons DC"
 where
  tys :: [Type]
tys = [(LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
n)), Type
elTy, (LitTy -> Type
LitTy (Integer -> LitTy
NumTy (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1)))]

-- | Create a vector of supplied elements
mkVecCons
  :: HasCallStack
  => DataCon
  -- ^ The Cons (:>) constructor
  -> Type
  -- ^ Element type
  -> Integer
  -- ^ Length of the vector
  -> Term
  -- ^ head of the vector
  -> Term
  -- ^ tail of the vector
  -> Term
mkVecCons :: HasCallStack => DataCon -> Type -> Integer -> Term -> Term -> Term
mkVecCons DataCon
consCon Type
resTy Integer
n Term
h Term
t
  | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
0 = [Char] -> Term
forall a. HasCallStack => [Char] -> a
error [Char]
"mkVecCons: n <= 0"
  | Bool
otherwise
  = case DataCon -> [Type] -> Maybe [Type]
dataConInstArgTys DataCon
consCon [LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
n), Type
resTy, LitTy -> Type
LitTy (Integer -> LitTy
NumTy (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1))] of
    Just (Type
consCoTy : [Type]
_) ->
      Term -> [Either Term Type] -> Term
mkApps (DataCon -> Term
Data DataCon
consCon) [ Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
n))
                            , Type -> Either Term Type
forall a b. b -> Either a b
Right Type
resTy
                            , Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1)))
                            , Term -> Either Term Type
forall a b. a -> Either a b
Left (Type -> Term
primCo Type
consCoTy)
                            , Term -> Either Term Type
forall a b. a -> Either a b
Left Term
h
                            , Term -> Either Term Type
forall a b. a -> Either a b
Left Term
t ]
    Maybe [Type]
_ -> [Char] -> Term
forall a. HasCallStack => [Char] -> a
error [Char]
"mkVecCons: failed to instantiate Cons DC"

-- | Create an empty vector
mkVecNil
  :: DataCon
  -- ^ The Nil constructor
  -> Type
  -- ^ The element type
  -> Term
mkVecNil :: DataCon -> Type -> Term
mkVecNil DataCon
nilCon Type
resTy = case DataCon -> [Type] -> Maybe [Type]
dataConInstArgTys DataCon
nilCon [LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
0), Type
resTy] of
  Just (Type
nilCoTy : [Type]
_) ->
    Term -> [Either Term Type] -> Term
mkApps (DataCon -> Term
Data DataCon
nilCon) [ Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
0))
                        , Type -> Either Term Type
forall a b. b -> Either a b
Right Type
resTy
                        , Term -> Either Term Type
forall a b. a -> Either a b
Left  (Type -> Term
primCo Type
nilCoTy) ]
  Maybe [Type]
_ -> [Char] -> Term
forall a. HasCallStack => [Char] -> a
error [Char]
"mkVecNil: failed to instantiate Nil DC"

-- | Replace an application of the @Clash.Sized.Vector.reverse@ primitive on
-- vectors of a known length @n@, by the fully unrolled recursive "definition"
-- of @Clash.Sized.Vector.reverse@
reduceReverse
  :: Integer
  -- ^ Length of the vector, must be positive
  -> Type
  -- ^ Element of type of the vector
  -> Term
  -- ^ The vector to reverse
  -> TransformContext
  -> NormalizeSession Term
reduceReverse :: Integer
-> Type -> Term -> TransformContext -> NormalizeSession Term
reduceReverse Integer
n Type
elTy Term
vArg (TransformContext InScopeSet
inScope0 Context
_ctx) = do
  tcm <- Getting TyConMap RewriteEnv TyConMap
-> RewriteMonad NormalizeState TyConMap
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting TyConMap RewriteEnv TyConMap
Getter RewriteEnv TyConMap
tcCache
  let ty = TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
vArg
  go tcm ty
 where
  go :: TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
tcm -> Just Type
ty') = TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
ty'
  go TyConMap
tcm (Type -> TypeView
tyView -> TyConApp TyConName
vecTcNm [Type]
_)
    | Just TyCon
vecTc <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
vecTcNm TyConMap
tcm
      , TyConName -> Text
forall a. Name a -> Text
nameOcc TyConName
vecTcNm Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"Clash.Sized.Vector.Vec"
    , [DataCon
nilCon, DataCon
consCon] <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
    = do
      uniqs0 <- Getting Supply (RewriteState extra) Supply
-> RewriteMonad extra Supply
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting Supply (RewriteState extra) Supply
forall extra (f :: Type -> Type).
Functor f =>
(Supply -> f Supply)
-> RewriteState extra -> f (RewriteState extra)
uniqSupply
      let (uniqs1,(vars,elems)) = second (second sconcat . NE.unzip)
                                $ extractElems uniqs0 inScope0 consCon elTy 'V' n vArg
          lbody = DataCon -> DataCon -> Type -> Integer -> [Term] -> Term
mkVec DataCon
nilCon DataCon
consCon Type
elTy Integer
n ([Term] -> [Term]
forall a. [a] -> [a]
reverse (NonEmpty Term -> [Term]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty Term
vars))
          lb    = [(Id, Term)] -> Term -> Term
Letrec (NonEmpty (Id, Term) -> [(Id, Term)]
forall a. NonEmpty a -> [a]
NE.init NonEmpty (Id, Term)
elems) Term
lbody
      uniqSupply Lens..= uniqs1
      changed lb
  go TyConMap
_ Type
ty = [Char] -> RewriteMonad extra Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> RewriteMonad extra Term)
-> [Char] -> RewriteMonad extra Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"reduceReverse: argument does not have a vector type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Type -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr Type
ty

-- | Replace an application of the @Clash.Sized.Vector.zipWith@ primitive on
-- vectors of a known length @n@, by the fully unrolled recursive "definition"
-- of @Clash.Sized.Vector.zipWith@
reduceZipWith
  :: PrimInfo -- ^ zipWith primitive info
  -> Integer  -- ^ Length of the vector(s)
  -> Type -- ^ Element type of the lhs of the function
  -> Type -- ^ Element type of the rhs of the function
  -> Type -- ^ Element type of the result of the function
  -> Term -- ^ The zipWith'd functions
  -> Term -- ^ The 1st vector argument
  -> Term -- ^ The 2nd vector argument
  -> TransformContext
  -> NormalizeSession Term
reduceZipWith :: PrimInfo
-> Integer
-> Type
-> Type
-> Type
-> Term
-> Term
-> Term
-> TransformContext
-> NormalizeSession Term
reduceZipWith PrimInfo
zipWithPrimInfo Integer
n Type
lhsElTy Type
rhsElTy Type
resElTy Term
fun Term
lhsArg Term
rhsArg TransformContext
_ctx = do
  tcm <- Getting TyConMap RewriteEnv TyConMap
-> RewriteMonad NormalizeState TyConMap
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting TyConMap RewriteEnv TyConMap
Getter RewriteEnv TyConMap
tcCache
  changed (go tcm (inferCoreTypeOf tcm lhsArg))
 where
  go :: TyConMap -> Type -> Term
go TyConMap
tcm (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
tcm -> Just Type
ty) = TyConMap -> Type -> Term
go TyConMap
tcm Type
ty
  go TyConMap
tcm (Type -> TypeView
tyView -> TyConApp TyConName
vecTcNm [Type]
_)
    | (Just TyCon
vecTc) <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
vecTcNm TyConMap
tcm
    , TyConName -> Text
forall a. Name a -> Text
nameOcc TyConName
vecTcNm Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"Clash.Sized.Vector.Vec"
    , [DataCon
nilCon, DataCon
consCon] <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
    = if Integer
n Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 then
        DataCon -> Type -> Term
mkVecNil DataCon
nilCon Type
resElTy
      else
        let
          (Term
a, Term
as) = DataCon -> Type -> Integer -> Term -> (Term, Term)
extractHeadTail DataCon
consCon Type
lhsElTy Integer
n Term
lhsArg
          (Term
b, Term
bs) = DataCon -> Type -> Integer -> Term -> (Term, Term)
extractHeadTail DataCon
consCon Type
rhsElTy Integer
n Term
rhsArg
          c :: Term
c = Term -> [Either Term Type] -> Term
mkApps Term
fun [Term -> Either Term Type
forall a b. a -> Either a b
Left Term
a, Term -> Either Term Type
forall a b. a -> Either a b
Left Term
b]
          cs :: Term
cs = Term -> [Either Term Type] -> Term
mkApps (PrimInfo -> Term
Prim PrimInfo
zipWithPrimInfo) [ Type -> Either Term Type
forall a b. b -> Either a b
Right Type
lhsElTy
                                             , Type -> Either Term Type
forall a b. b -> Either a b
Right Type
rhsElTy
                                             , Type -> Either Term Type
forall a b. b -> Either a b
Right Type
resElTy
                                             , Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy (Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1)))
                                             , Term -> Either Term Type
forall a b. a -> Either a b
Left Term
fun
                                             , Term -> Either Term Type
forall a b. a -> Either a b
Left Term
as
                                             , Term -> Either Term Type
forall a b. a -> Either a b
Left Term
bs ]
        in
          HasCallStack => DataCon -> Type -> Integer -> Term -> Term -> Term
DataCon -> Type -> Integer -> Term -> Term -> Term
mkVecCons DataCon
consCon Type
resElTy Integer
n Term
c Term
cs
  go TyConMap
_ Type
ty =
    [Char] -> Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> Term) -> [Char] -> Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [I.i|
      reduceZipWith: argument does not have a vector type:

        #{showPpr ty}
    |]

-- | Replace an application of the @Clash.Sized.Vector.map@ primitive on vectors
-- of a known length @n@, by the fully unrolled recursive "definition" of
-- @Clash.Sized.Vector.map@
reduceMap
  :: PrimInfo -- ^ map primitive info
  -> Integer  -- ^ Length of the vector
  -> Type -- ^ Argument type of the function
  -> Type -- ^ Result type of the function
  -> Term -- ^ The map'd function
  -> Term -- ^ The map'd over vector
  -> TransformContext
  -> NormalizeSession Term
reduceMap :: PrimInfo
-> Integer
-> Type
-> Type
-> Term
-> Term
-> TransformContext
-> NormalizeSession Term
reduceMap PrimInfo
mapPrimInfo Integer
n Type
argElTy Type
resElTy Term
fun Term
arg TransformContext
_ctx = do
    tcm <- Getting TyConMap RewriteEnv TyConMap
-> RewriteMonad NormalizeState TyConMap
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting TyConMap RewriteEnv TyConMap
Getter RewriteEnv TyConMap
tcCache
    let ty = TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
arg
    changed (go tcm ty)
  where
    go :: TyConMap -> Type -> Term
go TyConMap
tcm (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
tcm -> Just Type
ty') = TyConMap -> Type -> Term
go TyConMap
tcm Type
ty'
    go TyConMap
tcm (Type -> TypeView
tyView -> TyConApp TyConName
vecTcNm [Type]
_)
      | (Just TyCon
vecTc)     <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
vecTcNm TyConMap
tcm
      , TyConName -> Text
forall a. Name a -> Text
nameOcc TyConName
vecTcNm Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"Clash.Sized.Vector.Vec"
      , [DataCon
nilCon,DataCon
consCon] <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      = if Integer
n Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 then
          DataCon -> Type -> Term
mkVecNil DataCon
nilCon Type
argElTy
        else
          let
            nPredTy :: Either a Type
nPredTy = Type -> Either a Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy (Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1)))
            (Term
a, Term
as) = DataCon -> Type -> Integer -> Term -> (Term, Term)
extractHeadTail DataCon
consCon Type
argElTy Integer
n Term
arg
            b :: Term
b = Term -> [Either Term Type] -> Term
mkApps Term
fun [Term -> Either Term Type
forall a b. a -> Either a b
Left Term
a]
            bs :: Term
bs = Term -> [Either Term Type] -> Term
mkApps (PrimInfo -> Term
Prim PrimInfo
mapPrimInfo) [ Type -> Either Term Type
forall a b. b -> Either a b
Right Type
argElTy
                                           , Type -> Either Term Type
forall a b. b -> Either a b
Right Type
resElTy
                                           , Either Term Type
forall {a}. Either a Type
nPredTy
                                           , Term -> Either Term Type
forall a b. a -> Either a b
Left Term
fun
                                           , Term -> Either Term Type
forall a b. a -> Either a b
Left Term
as ]
          in
            HasCallStack => DataCon -> Type -> Integer -> Term -> Term -> Term
DataCon -> Type -> Integer -> Term -> Term -> Term
mkVecCons DataCon
consCon Type
resElTy Integer
n Term
b Term
bs
    go TyConMap
_ Type
ty =
      [Char] -> Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> Term) -> [Char] -> Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [I.i|
        reduceMap: argument does not have a vector type:

          #{showPpr ty}
      |]

-- | Replace an application of the @Clash.Sized.Vector.imap@ primitive on vectors
-- of a known length @n@, by the fully unrolled recursive "definition" of
-- @Clash.Sized.Vector.imap@
reduceImap
  :: Integer  -- ^ Length of the vector, must be positive
  -> Type -- ^ Argument type of the function
  -> Type -- ^ Result type of the function
  -> Term -- ^ Lenght of the vector (as a KnownNat)
  -> Term -- ^ The imap'd function
  -> Term -- ^ The imap'd over vector
  -> TransformContext
  -> NormalizeSession Term
reduceImap :: Integer
-> Type
-> Type
-> Term
-> Term
-> Term
-> TransformContext
-> NormalizeSession Term
reduceImap Integer
n Type
argElTy Type
resElTy Term
_kn Term
fun Term
arg (TransformContext InScopeSet
is0 Context
ctx) = do
    tcm <- Getting TyConMap RewriteEnv TyConMap
-> RewriteMonad NormalizeState TyConMap
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting TyConMap RewriteEnv TyConMap
Getter RewriteEnv TyConMap
tcCache
    let ty = TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
arg
    go tcm ty
  where
    go :: TyConMap -> Type -> NormalizeSession Term
go TyConMap
tcm (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
tcm -> Just Type
ty') = TyConMap -> Type -> NormalizeSession Term
go TyConMap
tcm Type
ty'
    go TyConMap
tcm (Type -> TypeView
tyView -> TyConApp TyConName
vecTcNm [Type]
_)
      | (Just TyCon
vecTc)     <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
vecTcNm TyConMap
tcm
      , TyConName -> Text
forall a. Name a -> Text
nameOcc TyConName
vecTcNm Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"Clash.Sized.Vector.Vec"
      , [DataCon
nilCon,DataCon
consCon] <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      = do
        uniqs0 <- Getting Supply (RewriteState NormalizeState) Supply
-> RewriteMonad NormalizeState Supply
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting Supply (RewriteState NormalizeState) Supply
forall extra (f :: Type -> Type).
Functor f =>
(Supply -> f Supply)
-> RewriteState extra -> f (RewriteState extra)
uniqSupply
        fun1 <- constantPropagation (TransformContext is0 (AppArg Nothing:ctx)) fun
        let is1 = InScopeSet -> [Id] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 (Term -> [Id]
collectTermIds Term
fun1)
            (uniqs1,nTv) = mkUniqSystemTyVar (uniqs0,is1) ("n",typeNatKind)
            (uniqs2,(vars,elems)) = second (second sconcat . NE.unzip)
                                  $ uncurry extractElems uniqs1 consCon argElTy 'I' n arg
            idxTcNm = TyConName -> Maybe TyConName -> TyConName
forall a. a -> Maybe a -> a
Maybe.fromMaybe ([Char] -> TyConName
forall a. HasCallStack => [Char] -> a
error [Char]
"reduceImap: failed to create Index TC") (Maybe TyConName -> TyConName) -> Maybe TyConName -> TyConName
forall a b. (a -> b) -> a -> b
$ do
              (Right idxTy:_,_) <- ([Either TyVar Type], Type) -> Maybe ([Either TyVar Type], Type)
forall a. a -> Maybe a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Type -> ([Either TyVar Type], Type)
splitFunForallTy (TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
fun))
              TyConApp nm _ <- pure (tyView idxTy)
              return nm
            -- fromInteger# :: KnownNat n => Integer -> Index n
            idxFromIntegerTy = TyVar -> Type -> Type
ForAllTy TyVar
nTv
                                        ((Type -> Type -> Type) -> Type -> [Type] -> Type
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Type -> Type -> Type
mkFunTy
                                               (TyConName -> [Type] -> Type
mkTyConApp TyConName
idxTcNm
                                                           [TyVar -> Type
VarTy TyVar
nTv])
                                               [Type
integerPrimTy,Type
integerPrimTy])
            idxFromInteger   = PrimInfo -> Term
Prim (Text
-> Type -> WorkInfo -> IsMultiPrim -> PrimUnfolding -> PrimInfo
PrimInfo Text
"Clash.Sized.Internal.Index.fromInteger#" Type
idxFromIntegerTy WorkInfo
WorkNever IsMultiPrim
SingleResult PrimUnfolding
NoUnfolding)
            idxs             = (Integer -> Term) -> [Integer] -> [Term]
forall a b. (a -> b) -> [a] -> [b]
map (Term -> Term -> Term
App (Term -> Term -> Term
App (Term -> Type -> Term
TyApp Term
idxFromInteger (LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
n)))
                                             (Literal -> Term
Literal (Integer -> Literal
IntegerLiteral (Integer -> Integer
forall a. Integral a => a -> Integer
toInteger Integer
n))))
                                   (Term -> Term) -> (Integer -> Term) -> Integer -> Term
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Literal -> Term
Literal (Literal -> Term) -> (Integer -> Literal) -> Integer -> Term
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Literal
IntegerLiteral (Integer -> Literal) -> (Integer -> Integer) -> Integer -> Literal
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Integer
forall a. Integral a => a -> Integer
toInteger) [Integer
0..(Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1)]
            funApps          = (Term -> Term -> Term) -> [Term] -> [Term] -> [Term]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Term
i Term
v -> Term -> Term -> Term
App (Term -> Term -> Term
App Term
fun1 Term
i) Term
v) [Term]
idxs (NonEmpty Term -> [Term]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty Term
vars)
            lbody            = DataCon -> DataCon -> Type -> Integer -> [Term] -> Term
mkVec DataCon
nilCon DataCon
consCon Type
resElTy Integer
n [Term]
funApps
            lb               = [(Id, Term)] -> Term -> Term
Letrec (NonEmpty (Id, Term) -> [(Id, Term)]
forall a. NonEmpty a -> [a]
NE.init NonEmpty (Id, Term)
elems) Term
lbody
        uniqSupply Lens..= uniqs2
        changed lb
    go TyConMap
_ Type
ty = [Char] -> NormalizeSession Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> NormalizeSession Term)
-> [Char] -> NormalizeSession Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"reduceImap: argument does not have a vector type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Type -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr Type
ty

-- | Replace an application of the @Clash.Sized.Vector.iterateI@ primitive on
-- vectors of a known length @n@, by the fully unrolled recursive "definition"
-- of @Clash.Sized.Vector.iterateI@
reduceIterateI
  :: Integer
  -- ^ Length of vector
  -> Type
  -- ^ Vector's element type
  -> Type
  -- ^ Vector's type
  -> Term
  -- ^ Length of the vector (as a KnownNat)
  -> Term
  -- ^ iterateI's HO-function argument
  -> Term
  -- ^ iterateI's start value
  -> TransformContext
  -> RewriteMonad NormalizeState Term
  -- ^ Fully unrolled definition
reduceIterateI :: Integer
-> Type
-> Type
-> Term
-> Term
-> Term
-> TransformContext
-> NormalizeSession Term
reduceIterateI Integer
n Type
aTy Type
vTy Term
_kn Term
f0 Term
a (TransformContext InScopeSet
is0 Context
ctx) = do
  tcm <- Getting TyConMap RewriteEnv TyConMap
-> RewriteMonad NormalizeState TyConMap
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting TyConMap RewriteEnv TyConMap
Getter RewriteEnv TyConMap
tcCache
  f1 <- constantPropagation (TransformContext is0 (AppArg Nothing:ctx)) f0

  -- Generate uniq ids for element assignments.
  uniqs0 <- Lens.use uniqSupply
  let
    is1 = InScopeSet -> [Id] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 (Term -> [Id]
collectTermIds Term
f1)
    ((uniqs1, _is2), elementIds) =
      mapAccumR
        mkUniqInternalId
        (uniqs0, is1)
        (zip (map (("el" <>) . showt) [1..n-1]) (repeat aTy))
  uniqSupply .= uniqs1

  let
    elems = (Term -> Term) -> [Term] -> [Term]
forall a b. (a -> b) -> [a] -> [b]
map (Term -> Term -> Term
App Term
f1) (Term
aTerm -> [Term] -> [Term]
forall a. a -> [a] -> [a]
:(Id -> Term) -> [Id] -> [Term]
forall a b. (a -> b) -> [a] -> [b]
map Id -> Term
Var [Id]
elementIds)
    vec = Term -> Maybe Term -> Term
forall a. a -> Maybe a -> a
Maybe.fromMaybe ([Char] -> Term
forall a. HasCallStack => [Char] -> a
error [Char]
"reduceIterateI: failed to create Vec DCs") (Maybe Term -> Term) -> Maybe Term -> Term
forall a b. (a -> b) -> a -> b
$ do
      TyConApp vecTcNm _ <- TypeView -> Maybe TypeView
forall a. a -> Maybe a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Type -> TypeView
tyView Type
vTy)
      vecTc <- UniqMap.lookup vecTcNm tcm
      [nilCon, consCon] <- pure (tyConDataCons vecTc)
      return (mkVec nilCon consCon aTy n (take (fromInteger n) (a:map Var elementIds)))

  -- Result:
  --   let
  --     el1 = f a
  --     el2 = f el1
  --     el3 = f el2
  --     ..
  --   in
  --     (a :> el1 :> el2 :> el3 :> ..)
  --
  changed (Letrec (zip elementIds elems) vec)

-- | Replace an application of the @Clash.Sized.Vector.traverse#@ primitive on
-- vectors of a known length @n@, by the fully unrolled recursive "definition"
-- of @Clash.Sized.Vector.traverse#@
reduceTraverse
  :: Integer  -- ^ Length of the vector, must be positive
  -> Type -- ^ Element type of the argument vector
  -> Type -- ^ The type of the applicative
  -> Type -- ^ Element type of the result vector
  -> Term -- ^ The @Applicative@ dictionary
  -> Term -- ^ The function to traverse with
  -> Term -- ^ The argument vector
  -> TransformContext
  -> NormalizeSession Term
reduceTraverse :: Integer
-> Type
-> Type
-> Type
-> Term
-> Term
-> Term
-> TransformContext
-> NormalizeSession Term
reduceTraverse Integer
n Type
aTy Type
fTy Type
bTy Term
dict Term
fun Term
arg (TransformContext InScopeSet
is0 Context
ctx) = do
    tcm <- Getting TyConMap RewriteEnv TyConMap
-> RewriteMonad NormalizeState TyConMap
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting TyConMap RewriteEnv TyConMap
Getter RewriteEnv TyConMap
tcCache
    case tyView (inferCoreTypeOf tcm dict) of
      TyConApp TyConName
apDictTcNm [Type]
_ ->
        let ty :: Type
ty = TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
arg
         in TyConMap -> TyConName -> Type -> NormalizeSession Term
forall {a}.
Uniquable a =>
TyConMap -> a -> Type -> NormalizeSession Term
go TyConMap
tcm TyConName
apDictTcNm Type
ty
      TypeView
t -> [Char] -> NormalizeSession Term
forall a. HasCallStack => [Char] -> a
error ([Char]
"reduceTraverse: expected a TC, but got: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> TypeView -> [Char]
forall a. Show a => a -> [Char]
show TypeView
t)
  where
    go :: TyConMap -> a -> Type -> NormalizeSession Term
go TyConMap
tcm a
apDictTcNm (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
tcm -> Just Type
ty') = TyConMap -> a -> Type -> NormalizeSession Term
go TyConMap
tcm a
apDictTcNm Type
ty'
    go TyConMap
tcm a
apDictTcNm (Type -> TypeView
tyView -> TyConApp TyConName
vecTcNm [Type]
_)
      | (Just TyCon
vecTc) <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
vecTcNm TyConMap
tcm
      , TyConName -> Text
forall a. Name a -> Text
nameOcc TyConName
vecTcNm Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"Clash.Sized.Vector.Vec"
      , [DataCon
nilCon,DataCon
consCon] <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      = (Maybe Term -> Term)
-> RewriteMonad NormalizeState (Maybe Term)
-> NormalizeSession Term
forall a b.
(a -> b)
-> RewriteMonad NormalizeState a -> RewriteMonad NormalizeState b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (Term -> Maybe Term -> Term
forall a. a -> Maybe a -> a
Maybe.fromMaybe ([Char] -> Term
forall a. HasCallStack => [Char] -> a
error [Char]
"reduceTraverse: failed to build")) (RewriteMonad NormalizeState (Maybe Term) -> NormalizeSession Term)
-> RewriteMonad NormalizeState (Maybe Term)
-> NormalizeSession Term
forall a b. (a -> b) -> a -> b
$ MaybeT (RewriteMonad NormalizeState) Term
-> RewriteMonad NormalizeState (Maybe Term)
forall (m :: Type -> Type) a. MaybeT m a -> m (Maybe a)
runMaybeT (MaybeT (RewriteMonad NormalizeState) Term
 -> RewriteMonad NormalizeState (Maybe Term))
-> MaybeT (RewriteMonad NormalizeState) Term
-> RewriteMonad NormalizeState (Maybe Term)
forall a b. (a -> b) -> a -> b
$ do
          uniqs0 <- Getting Supply (RewriteState NormalizeState) Supply
-> MaybeT (RewriteMonad NormalizeState) Supply
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting Supply (RewriteState NormalizeState) Supply
forall extra (f :: Type -> Type).
Functor f =>
(Supply -> f Supply)
-> RewriteState extra -> f (RewriteState extra)
uniqSupply
          fun1 <- lift (constantPropagation (TransformContext is0 (AppArg Nothing:ctx)) fun)
          let is1 = InScopeSet -> [Id] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 (Term -> [Id]
collectTermIds Term
fun1)
          apDictTc <- hoistMaybe (UniqMap.lookup apDictTcNm tcm)
          apDictCon <- hoistMaybe (Maybe.listToMaybe (tyConDataCons apDictTc))
          apDictIdTys <- hoistMaybe (dataConInstArgTys apDictCon [fTy])
          (uniqs1,apDictIds@[functorDictId,pureId,apId,_,_,_]) <- pure $
                mapAccumR mkUniqInternalId (uniqs0,is1)
                  (zipEqual ["functorDict","pure","ap","liftA2","apConstL","apConstR"]
                      apDictIdTys)

          TyConApp funcDictTcNm _ <- hoistMaybe (tyView <$> Maybe.listToMaybe apDictIdTys)
          funcDictTc <- hoistMaybe (UniqMap.lookup funcDictTcNm tcm)
          funcDictCon <- hoistMaybe (Maybe.listToMaybe (tyConDataCons funcDictTc))
          funcDictIdTys <- hoistMaybe (dataConInstArgTys funcDictCon [fTy])
          (uniqs2,funcDicIds@[fmapId,_]) <- pure $
                mapAccumR mkUniqInternalId uniqs1
                  (zipEqual ["fmap","fmapConst"] funcDictIdTys)

          let apPat    = DataCon -> [TyVar] -> [Id] -> Pat
DataPat DataCon
apDictCon   [] [Id]
apDictIds
              fnPat    = DataCon -> [TyVar] -> [Id] -> Pat
DataPat DataCon
funcDictCon [] [Id]
funcDicIds

              -- Extract the 'pure' function from the Applicative dictionary
              pureTy = Id -> Type
forall a. HasType a => a -> Type
coreTypeOf Id
pureId
              pureTm = Term -> Type -> [Alt] -> Term
Case Term
dict Type
pureTy [(Pat
apPat,Id -> Term
Var Id
pureId)]

              -- Extract the '<*>' function from the Applicative dictionary
              apTy   = Id -> Type
forall a. HasType a => a -> Type
coreTypeOf Id
apId
              apTm   = Term -> Type -> [Alt] -> Term
Case Term
dict Type
apTy [(Pat
apPat, Id -> Term
Var Id
apId)]

              -- Extract the Functor dictionary from the Applicative dictionary
              funcTy = Id -> Type
forall a. HasType a => a -> Type
coreTypeOf Id
functorDictId
              funcTm = Term -> Type -> [Alt] -> Term
Case Term
dict Type
funcTy
                                [(Pat
apPat,Id -> Term
Var Id
functorDictId)]

              -- Extract the 'fmap' function from the Functor dictionary
              fmapTy = Id -> Type
forall a. HasType a => a -> Type
coreTypeOf Id
fmapId
              fmapTm = Term -> Type -> [Alt] -> Term
Case (Id -> Term
Var Id
functorDictId) Type
fmapTy
                            [(Pat
fnPat, Id -> Term
Var Id
fmapId)]

              (uniqs3,(vars,elems)) = second (second sconcat . NE.unzip)
                                    $ uncurry extractElems uniqs2 consCon aTy 'T' n arg

              funApps = (Term -> Term) -> [Term] -> [Term]
forall a b. (a -> b) -> [a] -> [b]
map (Term
fun1 Term -> Term -> Term
`App`) (NonEmpty Term -> [Term]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty Term
vars)

              lbody   = TyConName
-> DataCon
-> DataCon
-> Term
-> Term
-> Term
-> Type
-> Integer
-> [Term]
-> Term
mkTravVec TyConName
vecTcNm DataCon
nilCon DataCon
consCon (Id -> Term
Var ([Id]
apDictIds[Id] -> Int -> Id
forall a. HasCallStack => [a] -> Int -> a
!!Int
1))
                                                        (Id -> Term
Var ([Id]
apDictIds[Id] -> Int -> Id
forall a. HasCallStack => [a] -> Int -> a
!!Int
2))
                                                        (Id -> Term
Var ([Id]
funcDicIds[Id] -> Int -> Id
forall a. HasCallStack => [a] -> Int -> a
!!Int
0))
                                                        Type
bTy Integer
n [Term]
funApps

              lb      = [(Id, Term)] -> Term -> Term
Letrec ([(([Id]
apDictIds[Id] -> Int -> Id
forall a. HasCallStack => [a] -> Int -> a
!!Int
0), Term
funcTm)
                                ,(([Id]
apDictIds[Id] -> Int -> Id
forall a. HasCallStack => [a] -> Int -> a
!!Int
1), Term
pureTm)
                                ,(([Id]
apDictIds[Id] -> Int -> Id
forall a. HasCallStack => [a] -> Int -> a
!!Int
2), Term
apTm)
                                ,(([Id]
funcDicIds[Id] -> Int -> Id
forall a. HasCallStack => [a] -> Int -> a
!!Int
0), Term
fmapTm)
                                ] [(Id, Term)] -> [(Id, Term)] -> [(Id, Term)]
forall a. [a] -> [a] -> [a]
++ NonEmpty (Id, Term) -> [(Id, Term)]
forall a. NonEmpty a -> [a]
NE.init NonEmpty (Id, Term)
elems) Term
lbody
          uniqSupply Lens..= uniqs3
          lift (changed lb)
    go TyConMap
_ a
_ Type
ty = [Char] -> NormalizeSession Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> NormalizeSession Term)
-> [Char] -> NormalizeSession Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"reduceTraverse: argument does not have a vector type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Type -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr Type
ty

-- | Create the traversable vector
--
-- e.g. for a length '2' input vector, we get
--
-- > (:>) <$> x0 <*> ((:>) <$> x1 <*> pure Nil)
mkTravVec :: TyConName -- ^ Vec tcon
          -> DataCon   -- ^ Nil con
          -> DataCon   -- ^ Cons con
          -> Term      -- ^ 'pure' term
          -> Term      -- ^ '<*>' term
          -> Term      -- ^ 'fmap' term
          -> Type      -- ^ 'b' ty
          -> Integer       -- ^ Length of the vector
          -> [Term]    -- ^ Elements of the vector
          -> Term
mkTravVec :: TyConName
-> DataCon
-> DataCon
-> Term
-> Term
-> Term
-> Type
-> Integer
-> [Term]
-> Term
mkTravVec TyConName
vecTc DataCon
nilCon DataCon
consCon Term
pureTm Term
apTm Term
fmapTm Type
bTy = Integer -> [Term] -> Term
go
  where
    go :: Integer -> [Term] -> Term
    go :: Integer -> [Term] -> Term
go Integer
_ [] = Term -> [Either Term Type] -> Term
mkApps Term
pureTm [Type -> Either Term Type
forall a b. b -> Either a b
Right (TyConName -> [Type] -> Type
mkTyConApp TyConName
vecTc [LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
0),Type
bTy])
                            ,Term -> Either Term Type
forall a b. a -> Either a b
Left  (Term -> [Either Term Type] -> Term
mkApps (DataCon -> Term
Data DataCon
nilCon)
                                           [Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
0))
                                           ,Type -> Either Term Type
forall a b. b -> Either a b
Right Type
bTy
                                           ,Term -> Either Term Type
forall a b. a -> Either a b
Left  (Type -> Term
primCo Type
nilCoTy)])]

    go Integer
n (Term
x:[Term]
xs) = Term -> [Either Term Type] -> Term
mkApps Term
apTm
      [Type -> Either Term Type
forall a b. b -> Either a b
Right (TyConName -> [Type] -> Type
mkTyConApp TyConName
vecTc [LitTy -> Type
LitTy (Integer -> LitTy
NumTy (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1)),Type
bTy])
      ,Type -> Either Term Type
forall a b. b -> Either a b
Right (TyConName -> [Type] -> Type
mkTyConApp TyConName
vecTc [LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
n),Type
bTy])
      ,Term -> Either Term Type
forall a b. a -> Either a b
Left (Term -> [Either Term Type] -> Term
mkApps Term
fmapTm [Type -> Either Term Type
forall a b. b -> Either a b
Right Type
bTy
                           ,Type -> Either Term Type
forall a b. b -> Either a b
Right (Type -> Type -> Type
mkFunTy (TyConName -> [Type] -> Type
mkTyConApp TyConName
vecTc [LitTy -> Type
LitTy (Integer -> LitTy
NumTy (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1)),Type
bTy])
                                           (TyConName -> [Type] -> Type
mkTyConApp TyConName
vecTc [LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
n),Type
bTy]))
                           ,Term -> Either Term Type
forall a b. a -> Either a b
Left  (Term -> [Either Term Type] -> Term
mkApps (DataCon -> Term
Data DataCon
consCon)
                                          [Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
n))
                                          ,Type -> Either Term Type
forall a b. b -> Either a b
Right Type
bTy
                                          ,Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1)))
                                          ,Term -> Either Term Type
forall a b. a -> Either a b
Left  (Type -> Term
primCo (Integer -> Type
consCoTy Integer
n))
                                          ])
                           ,Term -> Either Term Type
forall a b. a -> Either a b
Left  Term
x])
      ,Term -> Either Term Type
forall a b. a -> Either a b
Left (Integer -> [Term] -> Term
go (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1) [Term]
xs)]

    nilCoTy :: Type
nilCoTy = case DataCon -> [Type] -> Maybe [Type]
dataConInstArgTys DataCon
nilCon [(LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
0)), Type
bTy] of
                Just (Type
x:[Type]
_) -> Type
x
                Maybe [Type]
_ -> [Char] -> Type
forall a. HasCallStack => [Char] -> a
error [Char]
"impossible"

    consCoTy :: Integer -> Type
consCoTy Integer
n = case DataCon -> [Type] -> Maybe [Type]
dataConInstArgTys DataCon
consCon
                                        [(LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
n))
                                        ,Type
bTy
                                        ,(LitTy -> Type
LitTy (Integer -> LitTy
NumTy (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1)))] of
                   Just (Type
x:[Type]
_) -> Type
x
                   Maybe [Type]
_ -> [Char] -> Type
forall a. HasCallStack => [Char] -> a
error [Char]
"impossible"

-- | Replace an application of the @Clash.Sized.Vector.foldr@ primitive on
-- vectors of a known length @n@, by the fully unrolled recursive "definition"
-- of @Clash.Sized.Vector.foldr@
reduceFoldr
  :: PrimInfo
  -- ^ Primitive info for foldr blackbox
  -> Integer
  -- ^ Length of the vector
  -> Type
  -- ^ Element type of the argument vector
  -> Term
  -- ^ The function to fold with
  -> Term
  -- ^ The starting value
  -> Term
  -- ^ The argument vector
  -> TransformContext
  -> NormalizeSession Term
reduceFoldr :: PrimInfo
-> Integer
-> Type
-> Term
-> Term
-> Term
-> TransformContext
-> NormalizeSession Term
reduceFoldr PrimInfo
_ Integer
0 Type
_ Term
_ Term
start Term
_ TransformContext
_ = Term -> NormalizeSession Term
forall a extra. a -> RewriteMonad extra a
changed Term
start
reduceFoldr PrimInfo
foldrPrimInfo Integer
n Type
aTy Term
fun Term
start Term
arg TransformContext
_ctx = do
    tcm <- Getting TyConMap RewriteEnv TyConMap
-> RewriteMonad NormalizeState TyConMap
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting TyConMap RewriteEnv TyConMap
Getter RewriteEnv TyConMap
tcCache
    let ty = TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
arg
    changed (go tcm ty)
  where
    go :: TyConMap -> Type -> Term
go TyConMap
tcm (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
tcm -> Just Type
ty') = TyConMap -> Type -> Term
go TyConMap
tcm Type
ty'
    go TyConMap
tcm (Type -> TypeView
tyView -> TyConApp TyConName
vecTcNm [Type]
_)
      | TyConName -> Text
forall a. Name a -> Text
nameOcc TyConName
vecTcNm Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"Clash.Sized.Vector.Vec"
      , Just TyCon
vecTc <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
vecTcNm TyConMap
tcm
      , [DataCon
_nilCon, DataCon
consCon] <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      = let
          (Term
a, Term
as) = DataCon -> Type -> Integer -> Term -> (Term, Term)
extractHeadTail DataCon
consCon Type
aTy Integer
n Term
arg
          b :: Term
b = Term -> [Either Term Type] -> Term
mkApps (PrimInfo -> Term
Prim PrimInfo
foldrPrimInfo) [ Type -> Either Term Type
forall a b. b -> Either a b
Right Type
aTy
                                          , Type -> Either Term Type
forall a b. b -> Either a b
Right (TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
start)
                                          , Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy (Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1)))
                                          , Term -> Either Term Type
forall a b. a -> Either a b
Left Term
fun
                                          , Term -> Either Term Type
forall a b. a -> Either a b
Left Term
start
                                          , Term -> Either Term Type
forall a b. a -> Either a b
Left Term
as ]
        in
          Term -> [Either Term Type] -> Term
mkApps Term
fun [Term -> Either Term Type
forall a b. a -> Either a b
Left Term
a, Term -> Either Term Type
forall a b. a -> Either a b
Left Term
b]

    go TyConMap
_ Type
ty =
      [Char] -> Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> Term) -> [Char] -> Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [I.i|
        reduceFoldr: argument does not have a vector type:

          #{showPpr ty}
      |]

-- | Replace an application of the @Clash.Sized.Vector.fold@ primitive on
-- vectors of a known length @n@, by the fully unrolled recursive "definition"
-- of @Clash.Sized.Vector.fold@
reduceFold
  :: Integer
  -- ^ Length of the vector, must be positive
  -> Type
  -- ^ Element type of the argument vector
  -> Term
  -- ^ The function to fold with
  -> Term
  -- ^ The argument vector
  -> TransformContext
  -> NormalizeSession Term
reduceFold :: Integer
-> Type
-> Term
-> Term
-> TransformContext
-> NormalizeSession Term
reduceFold Integer
n Type
aTy Term
fun Term
arg (TransformContext InScopeSet
is0 Context
ctx) = do
    tcm <- Getting TyConMap RewriteEnv TyConMap
-> RewriteMonad NormalizeState TyConMap
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting TyConMap RewriteEnv TyConMap
Getter RewriteEnv TyConMap
tcCache
    let ty = TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
arg
    go tcm ty
  where
    go :: TyConMap -> Type -> NormalizeSession Term
go TyConMap
tcm (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
tcm -> Just Type
ty') = TyConMap -> Type -> NormalizeSession Term
go TyConMap
tcm Type
ty'
    go TyConMap
tcm (Type -> TypeView
tyView -> TyConApp TyConName
vecTcNm [Type]
_)
      | (Just TyCon
vecTc) <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
vecTcNm TyConMap
tcm
      , TyConName -> Text
forall a. Name a -> Text
nameOcc TyConName
vecTcNm Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"Clash.Sized.Vector.Vec"
      , [DataCon
_,DataCon
consCon]  <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      = do
        uniqs0 <- Getting Supply (RewriteState NormalizeState) Supply
-> RewriteMonad NormalizeState Supply
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting Supply (RewriteState NormalizeState) Supply
forall extra (f :: Type -> Type).
Functor f =>
(Supply -> f Supply)
-> RewriteState extra -> f (RewriteState extra)
uniqSupply
        fun1 <- constantPropagation (TransformContext is0 (AppArg Nothing:ctx)) fun
        let is1 = InScopeSet -> [Id] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 (Term -> [Id]
collectTermIds Term
fun1)
            (uniqs1,(vars,elems)) = second (second sconcat . NE.unzip)
                                  $ extractElems uniqs0 is1 consCon aTy 'F' n arg
            lbody            = Term -> [Term] -> Term
foldV Term
fun1 (NonEmpty Term -> [Term]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty Term
vars)
            lb               = [(Id, Term)] -> Term -> Term
Letrec (NonEmpty (Id, Term) -> [(Id, Term)]
forall a. NonEmpty a -> [a]
NE.init NonEmpty (Id, Term)
elems) Term
lbody
        uniqSupply Lens..= uniqs1
        changed lb
    go TyConMap
_ Type
ty = [Char] -> NormalizeSession Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> NormalizeSession Term)
-> [Char] -> NormalizeSession Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"reduceFold: argument does not have a vector type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Type -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr Type
ty

    foldV :: Term -> [Term] -> Term
foldV Term
_ [Term
a] = Term
a
    foldV Term
f [Term]
as  = let ([Term]
l,[Term]
r) = Int -> [Term] -> ([Term], [Term])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Term] -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [Term]
as Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) [Term]
as
                      lF :: Term
lF    = Term -> [Term] -> Term
foldV Term
f [Term]
l
                      rF :: Term
rF    = Term -> [Term] -> Term
foldV Term
f [Term]
r
                  in  Term -> [Either Term Type] -> Term
mkApps Term
f [Term -> Either Term Type
forall a b. a -> Either a b
Left Term
lF, Term -> Either Term Type
forall a b. a -> Either a b
Left Term
rF]

-- | Replace an application of the @Clash.Sized.Vector.dfold@ primitive on
-- vectors of a known length @n@, by the fully unrolled recursive "definition"
-- of @Clash.Sized.Vector.dfold@
reduceDFold
  :: Integer
  -- ^ Length of the vector
  -> Type
  -- ^ Element type of the argument vector
  -> Term
  -- ^ Length of the vector (as a KnownNat)
  -> Term
  -- ^ The motive
  -> Term
  -- ^ Function to fold with
  -> Term
  -- ^ Starting value
  -> Term
  -- ^ The vector to fold
  -> TransformContext
  -> NormalizeSession Term
reduceDFold :: Integer
-> Type
-> Term
-> Term
-> Term
-> Term
-> Term
-> TransformContext
-> NormalizeSession Term
reduceDFold Integer
0 Type
_   Term
_   Term
_       Term
_   Term
start Term
_   TransformContext
_ = Term -> NormalizeSession Term
forall a extra. a -> RewriteMonad extra a
changed Term
start
reduceDFold Integer
n Type
aTy Term
_kn Term
_motive Term
fun Term
start Term
arg (TransformContext InScopeSet
is0 Context
_ctx) = do
    tcm <- Getting TyConMap RewriteEnv TyConMap
-> RewriteMonad NormalizeState TyConMap
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting TyConMap RewriteEnv TyConMap
Getter RewriteEnv TyConMap
tcCache
    let ty = TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
arg
    go tcm ty
  where
    go :: TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
tcm -> Just Type
ty') = TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
ty'
    go TyConMap
tcm (Type -> TypeView
tyView -> TyConApp TyConName
vecTcNm [Type]
_)
      | (Just TyCon
vecTc) <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
vecTcNm TyConMap
tcm
      , TyConName -> Text
forall a. Name a -> Text
nameOcc TyConName
vecTcNm Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"Clash.Sized.Vector.Vec"
      , [DataCon
_,DataCon
consCon]  <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      = do
        uniqs0 <- Getting Supply (RewriteState extra) Supply
-> RewriteMonad extra Supply
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting Supply (RewriteState extra) Supply
forall extra (f :: Type -> Type).
Functor f =>
(Supply -> f Supply)
-> RewriteState extra -> f (RewriteState extra)
uniqSupply
        let is1 = InScopeSet -> [Id] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 (Term -> [Id]
collectTermIds Term
fun)
            -- TODO: Should 'constantPropagation' be used on 'fun'? It seems to
            -- TOOD: be used for every other function in this module.
            (uniqs1,(vars,elems)) = second (second sconcat . NE.unzip)
                                  $ extractElems uniqs0 is1 consCon aTy 'D' n arg
            snatDc = DataCon -> Maybe DataCon -> DataCon
forall a. a -> Maybe a -> a
Maybe.fromMaybe ([Char] -> DataCon
forall a. HasCallStack => [Char] -> a
error [Char]
"reduceDFold: faild to build SNat") (Maybe DataCon -> DataCon) -> Maybe DataCon -> DataCon
forall a b. (a -> b) -> a -> b
$ do
              (_ltv:Right snTy:_,_) <- ([Either TyVar Type], Type) -> Maybe ([Either TyVar Type], Type)
forall a. a -> Maybe a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Type -> ([Either TyVar Type], Type)
splitFunForallTy (TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
fun))
              (TyConApp snatTcNm _) <- pure (tyView snTy)
              snatTc <- UniqMap.lookup snatTcNm tcm
              Maybe.listToMaybe (tyConDataCons snatTc)
            lbody = (Integer -> Term) -> Integer -> [Term] -> Term
doFold (DataCon -> Integer -> Term
buildSNat DataCon
snatDc) (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1) (NonEmpty Term -> [Term]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty Term
vars)
            lb    = [(Id, Term)] -> Term -> Term
Letrec (NonEmpty (Id, Term) -> [(Id, Term)]
forall a. NonEmpty a -> [a]
NE.init NonEmpty (Id, Term)
elems) Term
lbody
        uniqSupply Lens..= uniqs1
        changed lb
    go TyConMap
_ Type
ty = [Char] -> RewriteMonad extra Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> RewriteMonad extra Term)
-> [Char] -> RewriteMonad extra Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"reduceDFold: argument does not have a vector type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Type -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr Type
ty

    doFold :: (Integer -> Term) -> Integer -> [Term] -> Term
doFold Integer -> Term
_    Integer
_ []     = Term
start
    doFold Integer -> Term
snDc Integer
k (Term
x:[Term]
xs) = Term -> [Either Term Type] -> Term
mkApps Term
fun
                                 [Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
k))
                                 ,Term -> Either Term Type
forall a b. a -> Either a b
Left (Integer -> Term
snDc Integer
k)
                                 ,Term -> Either Term Type
forall a b. a -> Either a b
Left Term
x
                                 ,Term -> Either Term Type
forall a b. a -> Either a b
Left ((Integer -> Term) -> Integer -> [Term] -> Term
doFold Integer -> Term
snDc (Integer
kInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1) [Term]
xs)
                                 ]

-- | Replace an application of the @Clash.Sized.Vector.head@ primitive on
-- vectors of a known length @n@, by a projection of the first element of a
-- vector.
reduceHead
  :: Integer  -- ^ Length of the vector, must be positive
  -> Type -- ^ Element type of the vector
  -> Term -- ^ The argument vector
  -> TransformContext
  -> NormalizeSession Term
reduceHead :: Integer
-> Type -> Term -> TransformContext -> NormalizeSession Term
reduceHead Integer
n Type
aTy Term
vArg (TransformContext InScopeSet
inScope Context
_ctx) = do
    tcm <- Getting TyConMap RewriteEnv TyConMap
-> RewriteMonad NormalizeState TyConMap
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting TyConMap RewriteEnv TyConMap
Getter RewriteEnv TyConMap
tcCache
    let ty = TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
vArg
    go tcm ty
  where
    go :: TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
tcm -> Just Type
ty') = TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
ty'
    go TyConMap
tcm (Type -> TypeView
tyView -> TyConApp TyConName
vecTcNm [Type]
_)
      | (Just TyCon
vecTc) <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
vecTcNm TyConMap
tcm
      , TyConName -> Text
forall a. Name a -> Text
nameOcc TyConName
vecTcNm Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"Clash.Sized.Vector.Vec"
      , [DataCon
_,DataCon
consCon]  <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      = do
        uniqs0 <- Getting Supply (RewriteState extra) Supply
-> RewriteMonad extra Supply
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting Supply (RewriteState extra) Supply
forall extra (f :: Type -> Type).
Functor f =>
(Supply -> f Supply)
-> RewriteState extra -> f (RewriteState extra)
uniqSupply
        let (uniqs1,(vars,elems)) = second (second sconcat . NE.unzip)
                                  $ extractElems uniqs0 inScope consCon aTy 'H' n vArg
            lb = [(Id, Term)] -> Term -> Term
Letrec [NonEmpty (Id, Term) -> (Id, Term)
forall a. NonEmpty a -> a
NE.head NonEmpty (Id, Term)
elems] (NonEmpty Term -> Term
forall a. NonEmpty a -> a
NE.head NonEmpty Term
vars)
        uniqSupply Lens..= uniqs1
        changed lb
    go TyConMap
_ Type
ty = [Char] -> RewriteMonad extra Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> RewriteMonad extra Term)
-> [Char] -> RewriteMonad extra Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"reduceHead: argument does not have a vector type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Type -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr Type
ty

-- | Replace an application of the @Clash.Sized.Vector.tail@ primitive on
-- vectors of a known length @n@, by a projection of the tail of a
-- vector.
reduceTail
  :: Integer  -- ^ Length of the vector, must be positive
  -> Type -- ^ Element type of the vector
  -> Term -- ^ The argument vector
  -> TransformContext
  -> NormalizeSession Term
reduceTail :: Integer
-> Type -> Term -> TransformContext -> NormalizeSession Term
reduceTail Integer
n Type
aTy Term
vArg (TransformContext InScopeSet
inScope Context
_ctx) = do
    tcm <- Getting TyConMap RewriteEnv TyConMap
-> RewriteMonad NormalizeState TyConMap
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting TyConMap RewriteEnv TyConMap
Getter RewriteEnv TyConMap
tcCache
    let ty = TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
vArg
    go tcm ty
  where
    go :: TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
tcm -> Just Type
ty') = TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
ty'
    go TyConMap
tcm (Type -> TypeView
tyView -> TyConApp TyConName
vecTcNm [Type]
_)
      | (Just TyCon
vecTc) <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
vecTcNm TyConMap
tcm
      , TyConName -> Text
forall a. Name a -> Text
nameOcc TyConName
vecTcNm Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"Clash.Sized.Vector.Vec"
      , [DataCon
_,DataCon
consCon]  <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      = do
        uniqs0 <- Getting Supply (RewriteState extra) Supply
-> RewriteMonad extra Supply
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting Supply (RewriteState extra) Supply
forall extra (f :: Type -> Type).
Functor f =>
(Supply -> f Supply)
-> RewriteState extra -> f (RewriteState extra)
uniqSupply
        let (uniqs1,(_,elems)) = second (second sconcat . NE.unzip)
                               $ extractElems uniqs0 inScope consCon aTy 'L' n vArg
            b@(tB,_)     = elems NE.!! 1
            lb           = [(Id, Term)] -> Term -> Term
Letrec [(Id, Term)
b] (Id -> Term
Var Id
tB)
        uniqSupply Lens..= uniqs1
        changed lb
    go TyConMap
_ Type
ty = [Char] -> RewriteMonad extra Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> RewriteMonad extra Term)
-> [Char] -> RewriteMonad extra Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"reduceTail: argument does not have a vector type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Type -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr Type
ty

-- | Replace an application of the @Clash.Sized.Vector.last@ primitive on
-- vectors of a known length @n@, by a projection of the last element of a
-- vector.
reduceLast
  :: Integer  -- ^ Length of the vector, must be positive
  -> Type -- ^ Element type of the vector
  -> Term -- ^ The argument vector
  -> TransformContext
  -> NormalizeSession Term
reduceLast :: Integer
-> Type -> Term -> TransformContext -> NormalizeSession Term
reduceLast Integer
n Type
aTy Term
vArg (TransformContext InScopeSet
inScope Context
_ctx) = do
    tcm <- Getting TyConMap RewriteEnv TyConMap
-> RewriteMonad NormalizeState TyConMap
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting TyConMap RewriteEnv TyConMap
Getter RewriteEnv TyConMap
tcCache
    let ty = TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
vArg
    go tcm ty
  where
    go :: TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
tcm -> Just Type
ty') = TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
ty'
    go TyConMap
tcm (Type -> TypeView
tyView -> TyConApp TyConName
vecTcNm [Type]
_)
      | (Just TyCon
vecTc) <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
vecTcNm TyConMap
tcm
      , TyConName -> Text
forall a. Name a -> Text
nameOcc TyConName
vecTcNm Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"Clash.Sized.Vector.Vec"
      , [DataCon
_,DataCon
consCon]  <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      = do
        uniqs0 <- Getting Supply (RewriteState extra) Supply
-> RewriteMonad extra Supply
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting Supply (RewriteState extra) Supply
forall extra (f :: Type -> Type).
Functor f =>
(Supply -> f Supply)
-> RewriteState extra -> f (RewriteState extra)
uniqSupply
        let (uniqs1,(_,elems)) = second NE.unzip
                               $ extractElems uniqs0 inScope consCon aTy 'L' n vArg
            (tB,_)       = NE.head (NE.last elems)
        uniqSupply Lens..= uniqs1
        case n of
         Integer
0 -> Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed (Term -> Type -> Term
TyApp (PrimInfo -> Term
Prim PrimInfo
NP.undefined) Type
aTy)
         Integer
_ -> Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed ([(Id, Term)] -> Term -> Term
Letrec (NonEmpty (Id, Term) -> [(Id, Term)]
forall a. NonEmpty a -> [a]
NE.init (NonEmpty (NonEmpty (Id, Term)) -> NonEmpty (Id, Term)
forall a. Semigroup a => NonEmpty a -> a
sconcat NonEmpty (NonEmpty (Id, Term))
elems)) (Id -> Term
Var Id
tB))
    go TyConMap
_ Type
ty = [Char] -> RewriteMonad extra Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> RewriteMonad extra Term)
-> [Char] -> RewriteMonad extra Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"reduceLast: argument does not have a vector type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Type -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr Type
ty

-- | Replace an application of the @Clash.Sized.Vector.init@ primitive on
-- vectors of a known length @n@, by a projection of the init of a
-- vector.
reduceInit
  :: PrimInfo -- ^ Primitive info for 'init'
  -> Integer  -- ^ Length of the vector
  -> Type -- ^ Element type of the vector
  -> Term -- ^ The argument vector
  -> TransformContext
  -> NormalizeSession Term
reduceInit :: PrimInfo
-> Integer
-> Type
-> Term
-> TransformContext
-> NormalizeSession Term
reduceInit PrimInfo
initPrimInfo Integer
n Type
aTy Term
vArg TransformContext
_ctx = do
  tcm <- Getting TyConMap RewriteEnv TyConMap
-> RewriteMonad NormalizeState TyConMap
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting TyConMap RewriteEnv TyConMap
Getter RewriteEnv TyConMap
tcCache
  let ty = TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
vArg
  changed (go tcm ty)
 where
  go :: TyConMap -> Type -> Term
go TyConMap
tcm (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
tcm -> Just Type
ty') = TyConMap -> Type -> Term
go TyConMap
tcm Type
ty'
  go TyConMap
tcm (Type -> TypeView
tyView -> TyConApp TyConName
vecTcNm [Type]
_)
    | (Just TyCon
vecTc) <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
vecTcNm TyConMap
tcm
    , TyConName -> Text
forall a. Name a -> Text
nameOcc TyConName
vecTcNm Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"Clash.Sized.Vector.Vec"
    , [DataCon
nilCon, DataCon
consCon]  <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
    = if Integer
n Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 then
        DataCon -> Type -> Term
mkVecNil DataCon
nilCon Type
aTy
      else
        let
          nPredTy :: Either a Type
nPredTy = Type -> Either a Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy (Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1)))
          (Term
a, Term
as0) = DataCon -> Type -> Integer -> Term -> (Term, Term)
extractHeadTail DataCon
consCon Type
aTy (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
1) Term
vArg
          as1 :: Term
as1 = Term -> [Either Term Type] -> Term
mkApps (PrimInfo -> Term
Prim PrimInfo
initPrimInfo) [Either Term Type
forall {a}. Either a Type
nPredTy, Type -> Either Term Type
forall a b. b -> Either a b
Right Type
aTy, Term -> Either Term Type
forall a b. a -> Either a b
Left Term
as0]
        in
          HasCallStack => DataCon -> Type -> Integer -> Term -> Term -> Term
DataCon -> Type -> Integer -> Term -> Term -> Term
mkVecCons DataCon
consCon Type
aTy Integer
n Term
a Term
as1

  go TyConMap
_ Type
ty =
    [Char] -> Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> Term) -> [Char] -> Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [I.i|
      reduceInit: argument does not have a vector type:

        #{showPpr ty}
    |]

-- | Replace an application of the @Clash.Sized.Vector.(++)@ primitive on
-- vectors of a known length @n@, by the fully unrolled recursive "definition"
-- of @Clash.Sized.Vector.(++)@
reduceAppend
  :: Integer  -- ^ Length of the LHS arg
  -> Integer  -- ^ Lenght of the RHS arg
  -> Type -- ^ Element type of the vectors
  -> Term -- ^ The LHS argument
  -> Term -- ^ The RHS argument
  -> TransformContext
  -> NormalizeSession Term
reduceAppend :: Integer
-> Integer
-> Type
-> Term
-> Term
-> TransformContext
-> NormalizeSession Term
reduceAppend Integer
0 Integer
_ Type
_   Term
_    Term
rArg TransformContext
_ = Term -> NormalizeSession Term
forall a extra. a -> RewriteMonad extra a
changed Term
rArg
reduceAppend Integer
_ Integer
0 Type
_   Term
lArg Term
_    TransformContext
_ = Term -> NormalizeSession Term
forall a extra. a -> RewriteMonad extra a
changed Term
lArg
reduceAppend Integer
n Integer
m Type
aTy Term
lArg Term
rArg (TransformContext InScopeSet
inScope Context
_ctx) = do
    tcm <- Getting TyConMap RewriteEnv TyConMap
-> RewriteMonad NormalizeState TyConMap
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting TyConMap RewriteEnv TyConMap
Getter RewriteEnv TyConMap
tcCache
    let ty = TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
lArg
    go tcm ty
  where
    go :: TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
tcm -> Just Type
ty') = TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
ty'
    go TyConMap
tcm (Type -> TypeView
tyView -> TyConApp TyConName
vecTcNm [Type]
_)
      | (Just TyCon
vecTc) <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
vecTcNm TyConMap
tcm
      , TyConName -> Text
forall a. Name a -> Text
nameOcc TyConName
vecTcNm Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"Clash.Sized.Vector.Vec"
      , [DataCon
_,DataCon
consCon]  <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      = do uniqs0 <- Getting Supply (RewriteState extra) Supply
-> RewriteMonad extra Supply
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting Supply (RewriteState extra) Supply
forall extra (f :: Type -> Type).
Functor f =>
(Supply -> f Supply)
-> RewriteState extra -> f (RewriteState extra)
uniqSupply
           let (uniqs1,(vars,elems)) = second (second sconcat . NE.unzip)
                                     $ extractElems uniqs0 inScope consCon aTy
                                         'C' n lArg
               lbody        = DataCon -> Type -> Term -> Integer -> [Term] -> Term
appendToVec DataCon
consCon Type
aTy Term
rArg (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
m) (NonEmpty Term -> [Term]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty Term
vars)
               lb           = [(Id, Term)] -> Term -> Term
Letrec (NonEmpty (Id, Term) -> [(Id, Term)]
forall a. NonEmpty a -> [a]
NE.init NonEmpty (Id, Term)
elems) Term
lbody
           uniqSupply Lens..= uniqs1
           changed lb
    go TyConMap
_ Type
ty = [Char] -> RewriteMonad extra Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> RewriteMonad extra Term)
-> [Char] -> RewriteMonad extra Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"reduceAppend: argument does not have a vector type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Type -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr Type
ty

-- | Replace an application of the @Clash.Sized.Vector.unconcat@ primitive on
-- vectors of a known length @n@, by the fully unrolled recursive "definition"
-- of @Clash.Sized.Vector.unconcat@
reduceUnconcat :: PrimInfo -- ^ Unconcat primitive info
               -> Integer  -- ^ Length of the result vector
               -> Integer  -- ^ Length of the elements of the result vector
               -> Type -- ^ Element type
               -> Term -- ^ Length of the result vector (as a KnownNat)
               -> Term -- ^ SNat "Length of the elements of the result vector"
               -> Term -- ^ Argument vector
               -> TransformContext
               -> NormalizeSession Term
reduceUnconcat :: PrimInfo
-> Integer
-> Integer
-> Type
-> Term
-> Term
-> Term
-> TransformContext
-> NormalizeSession Term
reduceUnconcat PrimInfo
unconcatPrimInfo Integer
n Integer
m Type
aTy Term
_kn Term
sm Term
arg (TransformContext InScopeSet
inScope Context
_ctx) = do
    tcm <- Getting TyConMap RewriteEnv TyConMap
-> RewriteMonad NormalizeState TyConMap
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting TyConMap RewriteEnv TyConMap
Getter RewriteEnv TyConMap
tcCache
    let ty = TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
arg
    go tcm ty
  where
    go :: TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
tcm -> Just Type
ty') = TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
ty'
    go TyConMap
tcm (Type -> TypeView
tyView -> TyConApp TyConName
vecTcNm [Type]
_)
      | (Just TyCon
vecTc)     <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
vecTcNm TyConMap
tcm
      , TyConName -> Text
forall a. Name a -> Text
nameOcc TyConName
vecTcNm Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"Clash.Sized.Vector.Vec"
      , [DataCon
nilCon,DataCon
consCon] <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      , let innerVecTy :: Type
innerVecTy = TyConName -> [Type] -> Type
mkTyConApp TyConName
vecTcNm [LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
m), Type
aTy]
      = if Integer
n Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 then
          Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed (DataCon -> Type -> Term
mkVecNil DataCon
nilCon Type
innerVecTy)
        else if Integer
m Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 then do
          let
            nilVec :: Term
nilVec = DataCon -> Type -> Term
mkVecNil DataCon
nilCon Type
aTy
            retVec :: Term
retVec = DataCon -> DataCon -> Type -> Integer -> [Term] -> Term
mkVec DataCon
nilCon DataCon
consCon Type
innerVecTy Integer
n (Int -> Term -> [Term]
forall a. Int -> a -> [a]
replicate (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
n) Term
nilVec)
          Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed Term
retVec
        else do
          uniqs0 <- Getting Supply (RewriteState extra) Supply
-> RewriteMonad extra Supply
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting Supply (RewriteState extra) Supply
forall extra (f :: Type -> Type).
Functor f =>
(Supply -> f Supply)
-> RewriteState extra -> f (RewriteState extra)
uniqSupply
          let
            (uniqs1,(vars,headsAndTails)) =
              second (second sconcat . NE.unzip)
                     (extractElems uniqs0 inScope consCon aTy 'U' (n*m) arg)
            -- Build a vector out of the first m elements
            mvec = DataCon -> DataCon -> Type -> Integer -> [Term] -> Term
mkVec DataCon
nilCon DataCon
consCon Type
aTy Integer
m (Int -> NonEmpty Term -> [Term]
forall a. Int -> NonEmpty a -> [a]
NE.take (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
m) NonEmpty Term
vars)
            -- Get the vector representing the next ((n-1)*m) elements
            -- N.B. `extractElems (xs :: Vec 2 a)` creates:
            -- x0  = head xs
            -- xs0 = tail xs
            -- x1  = head xs0
            -- xs1 = tail xs0
            (lbs,nextVec) = case NE.splitAt ((2*fromInteger m)-1) headsAndTails of
                              ([(Id, Term)]
xs,(Id, Term)
y:[(Id, Term)]
_) -> ([(Id, Term)]
xs,(Id, Term)
y)
                              ([(Id, Term)], [(Id, Term)])
_ -> [Char] -> ([(Id, Term)], (Id, Term))
forall a. HasCallStack => [Char] -> a
error [Char]
"impossible"
            -- recursively call unconcat
            nextUnconcat = Term -> [Either Term Type] -> Term
mkApps (PrimInfo -> Term
Prim PrimInfo
unconcatPrimInfo)
                                  [ Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1)))
                                  , Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
m))
                                  , Type -> Either Term Type
forall a b. b -> Either a b
Right Type
aTy
                                  , Term -> Either Term Type
forall a b. a -> Either a b
Left (Literal -> Term
Literal (Integer -> Literal
NaturalLiteral (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1)))
                                  , Term -> Either Term Type
forall a b. a -> Either a b
Left Term
sm
                                  , Term -> Either Term Type
forall a b. a -> Either a b
Left ((Id, Term) -> Term
forall a b. (a, b) -> b
snd (Id, Term)
nextVec)
                                  ]
            -- let (mvec,nextVec) = splitAt sm arg
            -- in Cons mvec (unconcat sm nextVec)
            lBody = HasCallStack => DataCon -> Type -> Integer -> Term -> Term -> Term
DataCon -> Type -> Integer -> Term -> Term -> Term
mkVecCons DataCon
consCon Type
innerVecTy Integer
n Term
mvec Term
nextUnconcat
            lb = [(Id, Term)] -> Term -> Term
Letrec [(Id, Term)]
lbs Term
lBody

          uniqSupply Lens..= uniqs1
          changed lb
    go TyConMap
_ Type
ty = [Char] -> RewriteMonad extra Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> RewriteMonad extra Term)
-> [Char] -> RewriteMonad extra Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"reduceUnconcat: argument does not have a vector type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Type -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr Type
ty

-- | Replace an application of the @Clash.Sized.Vector.transpose@ primitive on
-- vectors of a known length @n@, by the fully unrolled recursive "definition"
-- of @Clash.Sized.Vector.transpose@
reduceTranspose :: Integer  -- ^ Length of the result vector
                -> Integer  -- ^ Length of the elements of the result vector
                -> Type -- ^ Element type
                -> Term -- ^ Lenght of the result vector (as a KnownNat)
                -> Term -- ^ Argument vector
                -> TransformContext
                -> NormalizeSession Term
reduceTranspose :: Integer
-> Integer
-> Type
-> Term
-> Term
-> TransformContext
-> NormalizeSession Term
reduceTranspose Integer
n Integer
0 Type
aTy Term
_kn Term
arg TransformContext
_ctx = do
    tcm <- Getting TyConMap RewriteEnv TyConMap
-> RewriteMonad NormalizeState TyConMap
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting TyConMap RewriteEnv TyConMap
Getter RewriteEnv TyConMap
tcCache
    let ty = TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
arg
    go tcm ty
  where
    go :: TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
tcm -> Just Type
ty') = TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
ty'
    go TyConMap
tcm (Type -> TypeView
tyView -> TyConApp TyConName
vecTcNm [Type]
_)
      | (Just TyCon
vecTc)     <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
vecTcNm TyConMap
tcm
      , TyConName -> Text
forall a. Name a -> Text
nameOcc TyConName
vecTcNm Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"Clash.Sized.Vector.Vec"
      , [DataCon
nilCon,DataCon
consCon] <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      = let nilVec :: Term
nilVec           = DataCon -> DataCon -> Type -> Integer -> [Term] -> Term
mkVec DataCon
nilCon DataCon
consCon Type
aTy Integer
0 []
            innerVecTy :: Type
innerVecTy       = TyConName -> [Type] -> Type
mkTyConApp TyConName
vecTcNm [LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
0), Type
aTy]
            retVec :: Term
retVec           = DataCon -> DataCon -> Type -> Integer -> [Term] -> Term
mkVec DataCon
nilCon DataCon
consCon Type
innerVecTy Integer
n (Int -> Term -> [Term]
forall a. Int -> a -> [a]
replicate (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
n) Term
nilVec)
        in  Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed Term
retVec
    go TyConMap
_ Type
ty = [Char] -> RewriteMonad extra Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> RewriteMonad extra Term)
-> [Char] -> RewriteMonad extra Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"reduceTranspose: argument does not have a vector type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Type -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr Type
ty

reduceTranspose Integer
_ Integer
_ Type
_ Term
_ Term
_ TransformContext
_ = [Char] -> NormalizeSession Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> NormalizeSession Term)
-> [Char] -> NormalizeSession Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"reduceTranspose: unimplemented"

reduceReplicate :: Integer
                -> Type
                -> Type
                -> Term
                -> Term
                -> TransformContext
                -> NormalizeSession Term
reduceReplicate :: Integer
-> Type
-> Type
-> Term
-> Term
-> TransformContext
-> NormalizeSession Term
reduceReplicate Integer
n Type
aTy Type
eTy Term
_sn Term
arg TransformContext
_ctx = do
    tcm <- Getting TyConMap RewriteEnv TyConMap
-> RewriteMonad NormalizeState TyConMap
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting TyConMap RewriteEnv TyConMap
Getter RewriteEnv TyConMap
tcCache
    go tcm eTy
  where
    go :: TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
tcm -> Just Type
ty') = TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
ty'
    go TyConMap
tcm (Type -> TypeView
tyView -> TyConApp TyConName
vecTcNm [Type]
_)
      | (Just TyCon
vecTc)     <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
vecTcNm TyConMap
tcm
      , TyConName -> Text
forall a. Name a -> Text
nameOcc TyConName
vecTcNm Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"Clash.Sized.Vector.Vec"
      , [DataCon
nilCon,DataCon
consCon] <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      = let retVec :: Term
retVec = DataCon -> DataCon -> Type -> Integer -> [Term] -> Term
mkVec DataCon
nilCon DataCon
consCon Type
aTy Integer
n (Int -> Term -> [Term]
forall a. Int -> a -> [a]
replicate (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
n) Term
arg)
        in  Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed Term
retVec
    go TyConMap
_ Type
ty = [Char] -> RewriteMonad extra Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> RewriteMonad extra Term)
-> [Char] -> RewriteMonad extra Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"reduceReplicate: argument does not have a vector type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Type -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr Type
ty

-- TODO: Take a shortcut when given index is a literal. Right now, this function
-- TODO: simply creates a case statement for every element in the vector, which
-- TODO: Clash will eliminate one-by-one if the index turned out to be literal.
-- TODO: It would of course be best to not create the cases in the first place!
reduceReplace_int
  :: Integer
  -- ^ Size of vector, must be positive
  -> Type
  -- ^ Type of vector element
  -> Type
  -- ^ Type of vector
  -> Term
  -- ^ Size of vector (as a KnownNat)
  -> Term
  -- ^ Vector
  -> Term
  -- ^ Index
  -> Term
  -- ^ Element
  -> TransformContext
  -> NormalizeSession Term
reduceReplace_int :: Integer
-> Type
-> Type
-> Term
-> Term
-> Term
-> Term
-> TransformContext
-> NormalizeSession Term
reduceReplace_int Integer
n Type
aTy Type
vTy Term
_kn Term
v Term
i Term
newA (TransformContext InScopeSet
is0 Context
_ctx) = do
  tcm <- Getting TyConMap RewriteEnv TyConMap
-> RewriteMonad NormalizeState TyConMap
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting TyConMap RewriteEnv TyConMap
Getter RewriteEnv TyConMap
tcCache
  go tcm vTy
 where
  -- Basically creates:
  --
  -- case eqInt i0 curI  of
  --   True -> newA
  --   _    -> oldA
  --
  -- where:
  --
  --   - curI is the index of the current element, which we statically know
  --   - i0 is the index given to replace_int
  --   - newA is the element given to replace_int as a replacement for..
  --   - oldA; an element at index curI
  --
  replace_intElement
    :: TyConMap
    -- ^ TyCon map
    -> DataCon
    -- Int datacon
    -> Type
    -- Int type
    -> Term
    -- ^ Element in vector
    -> Integer
    -- ^
    -> Term
  replace_intElement :: TyConMap -> DataCon -> Type -> Term -> Integer -> Term
replace_intElement TyConMap
tcm DataCon
iDc Type
iTy Term
oldA Integer
elIndex = Term
case0
   where
    case0 :: Term
case0 = Term -> Maybe Term -> Term
forall a. a -> Maybe a -> a
Maybe.fromMaybe ([Char] -> Term
forall a. HasCallStack => [Char] -> a
error [Char]
"replace_intElement: faild to build Truce DC") (Maybe Term -> Term) -> Maybe Term -> Term
forall a b. (a -> b) -> a -> b
$ do
      boolTc <- Unique -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup (Unique -> Unique
fromGhcUnique Unique
boolTyConKey) TyConMap
tcm
      [_,trueDc] <- pure (tyConDataCons boolTc)
      let eqInt = Type -> Type -> Term
eqIntPrim Type
iTy (TyConName -> [Type] -> Type
mkTyConApp (TyCon -> TyConName
tyConName TyCon
boolTc) [])
      return (Case (mkApps eqInt [ Left i
                                 , Left (mkApps (Data iDc)
                                                [Left (Literal (IntLiteral elIndex))])
                                 ])
                   aTy
                   [ (DefaultPat, oldA)
                   , (DataPat trueDc [] [], newA)
                   ])

  -- Equality on lifted Int that returns a Bool
  eqIntPrim
    :: Type
    -> Type
    -> Term
  eqIntPrim :: Type -> Type -> Term
eqIntPrim Type
intTy Type
boolTy =
    PrimInfo -> Term
Prim (Text
-> Type -> WorkInfo -> IsMultiPrim -> PrimUnfolding -> PrimInfo
PrimInfo
           Text
"GHC.Classes.eqInt"
           (Type -> Type -> Type
mkFunTy Type
intTy (Type -> Type -> Type
mkFunTy Type
intTy Type
boolTy))
           WorkInfo
WorkVariable
           IsMultiPrim
SingleResult
           PrimUnfolding
NoUnfolding)

  go :: TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
tcm -> Just Type
ty') = TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
ty'
  go TyConMap
tcm (Type -> TypeView
tyView -> TyConApp TyConName
vecTcNm [Type]
_)
    | (Just TyCon
vecTc)     <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
vecTcNm TyConMap
tcm
      , TyConName -> Text
forall a. Name a -> Text
nameOcc TyConName
vecTcNm Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"Clash.Sized.Vector.Vec"
    , [DataCon
nilCon,DataCon
consCon] <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
    = do
      -- Get data constructors of 'Int'
      uniqs0                   <- Getting Supply (RewriteState extra) Supply
-> RewriteMonad extra Supply
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting Supply (RewriteState extra) Supply
forall extra (f :: Type -> Type).
Functor f =>
(Supply -> f Supply)
-> RewriteState extra -> f (RewriteState extra)
uniqSupply
      let iTy                   = TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
i
          iDc = DataCon -> Maybe DataCon -> DataCon
forall a. a -> Maybe a -> a
Maybe.fromMaybe ([Char] -> DataCon
forall a. HasCallStack => [Char] -> a
error [Char]
"replace_intElement: faild to build Int DC") (Maybe DataCon -> DataCon) -> Maybe DataCon -> DataCon
forall a b. (a -> b) -> a -> b
$ do
            (TyConApp iTcNm _) <- TypeView -> Maybe TypeView
forall a. a -> Maybe a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Type -> TypeView
tyView Type
iTy)
            iTc <- UniqMap.lookup iTcNm tcm
            Maybe.listToMaybe (tyConDataCons iTc)

      -- Get elements from vector
          (uniqs1,(vars,elems)) = second (second sconcat . NE.unzip)
                                $ extractElems
                                    uniqs0
                                    is0
                                    consCon
                                    aTy
                                    'I'
                                    n
                                    v

      -- Replace every element with (if i == elIndex then newA else oldA)
      let replacedEls = (Term -> Integer -> Term) -> [Term] -> [Integer] -> [Term]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (TyConMap -> DataCon -> Type -> Term -> Integer -> Term
replace_intElement TyConMap
tcm DataCon
iDc Type
iTy) (NonEmpty Term -> [Term]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty Term
vars) [Integer
0..]
          lbody       = DataCon -> DataCon -> Type -> Integer -> [Term] -> Term
mkVec DataCon
nilCon DataCon
consCon Type
aTy Integer
n [Term]
replacedEls
          lb          = [(Id, Term)] -> Term -> Term
Letrec (NonEmpty (Id, Term) -> [(Id, Term)]
forall a. NonEmpty a -> [a]
NE.init NonEmpty (Id, Term)
elems) Term
lbody
      uniqSupply Lens..= uniqs1
      changed lb
  go TyConMap
_ Type
ty = [Char] -> RewriteMonad extra Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> RewriteMonad extra Term)
-> [Char] -> RewriteMonad extra Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"reduceReplace_int: argument does not have "
                                [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"a vector type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Type -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr Type
ty

-- TODO: Take a shortcut when given index is a literal. Right now, this function
-- TODO: simply creates a case statement for every element in the vector, which
-- TODO: Clash will eliminate one-by-one if the index turned out to be literal.
-- TODO: It would of course be best to not create the cases in the first place!
reduceIndex_int
  :: Integer
  -- ^ Size of vector, must be positive
  -> Type
  -- ^ Type of vector element
  -> Term
  -- ^ Size of vector (as a KnownNat)
  -> Term
  -- ^ Vector
  -> Term
  -- ^ Index
  -> TransformContext
  -> NormalizeSession Term
reduceIndex_int :: Integer
-> Type
-> Term
-> Term
-> Term
-> TransformContext
-> NormalizeSession Term
reduceIndex_int Integer
n Type
aTy Term
_kn Term
v Term
i (TransformContext InScopeSet
is0 Context
_ctx) = do
  tcm <- Getting TyConMap RewriteEnv TyConMap
-> RewriteMonad NormalizeState TyConMap
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting TyConMap RewriteEnv TyConMap
Getter RewriteEnv TyConMap
tcCache
  let vTy = TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
v
  go tcm vTy
 where
  -- Basically creates:
  --
  -- case eqInt i0 curI of
  --   True -> curA
  --   _    -> next
  --
  -- where:
  --
  --   - curI is the index of the current element, which we statically know
  --   - i0 is the index given to index_int
  --   - curA is the element at index curI
  --   - next; the value if the current index is not equal to index argument
  --
  index_intElement
    :: TyConMap
    -- ^ TyCon map
    -> DataCon
    -- Int datacon
    -> Type
    -- Int type
    -> (Term, Integer)
    -- ^ Element in the vector, and its corresponding index
    -> Term
    -- ^ The rest
    -> Term
  index_intElement :: TyConMap -> DataCon -> Type -> (Term, Integer) -> Term -> Term
index_intElement TyConMap
tcm DataCon
iDc Type
iTy (Term
cur,Integer
elIndex) Term
next = Term
case0
   where
    case0 :: Term
case0 = Term -> Maybe Term -> Term
forall a. a -> Maybe a -> a
Maybe.fromMaybe ([Char] -> Term
forall a. HasCallStack => [Char] -> a
error [Char]
"reduceIndex_int: faild to build True DC") (Maybe Term -> Term) -> Maybe Term -> Term
forall a b. (a -> b) -> a -> b
$ do
      boolTc <- Unique -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup (Unique -> Unique
fromGhcUnique Unique
boolTyConKey) TyConMap
tcm
      [_,trueDc] <- pure (tyConDataCons boolTc)
      let eqInt = Type -> Type -> Term
eqIntPrim Type
iTy (TyConName -> [Type] -> Type
mkTyConApp (TyCon -> TyConName
tyConName TyCon
boolTc) [])
      return (Case (mkApps eqInt [ Left i
                                 , Left (mkApps (Data iDc)
                                        [Left (Literal (IntLiteral elIndex))])
                                 ])
                   aTy
                   [ (DefaultPat, next)
                   , (DataPat trueDc [] [], cur)
                   ])

  -- Equality on lifted Int that returns a Bool
  eqIntPrim
    :: Type
    -> Type
    -> Term
  eqIntPrim :: Type -> Type -> Term
eqIntPrim Type
intTy Type
boolTy =
    PrimInfo -> Term
Prim ( Text
-> Type -> WorkInfo -> IsMultiPrim -> PrimUnfolding -> PrimInfo
PrimInfo
            Text
"GHC.Classes.eqInt"
            (Type -> Type -> Type
mkFunTy Type
intTy (Type -> Type -> Type
mkFunTy Type
intTy Type
boolTy))
            WorkInfo
WorkVariable
            IsMultiPrim
SingleResult
            PrimUnfolding
NoUnfolding)

  go :: TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
tcm -> Just Type
ty') = TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
ty'
  go TyConMap
tcm (Type -> TypeView
tyView -> TyConApp TyConName
vecTcNm [Type]
_)
    | (Just TyCon
vecTc)     <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
vecTcNm TyConMap
tcm
    , TyConName -> Text
forall a. Name a -> Text
nameOcc TyConName
vecTcNm Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"Clash.Sized.Vector.Vec"
    , [DataCon
_nilCon,DataCon
consCon] <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
    = do
      -- Get data constructors of 'Int'
      uniqs0                   <- Getting Supply (RewriteState extra) Supply
-> RewriteMonad extra Supply
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting Supply (RewriteState extra) Supply
forall extra (f :: Type -> Type).
Functor f =>
(Supply -> f Supply)
-> RewriteState extra -> f (RewriteState extra)
uniqSupply
      let iTy                   = TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
i
          iDc = DataCon -> Maybe DataCon -> DataCon
forall a. a -> Maybe a -> a
Maybe.fromMaybe ([Char] -> DataCon
forall a. HasCallStack => [Char] -> a
error [Char]
"reduceIndex_int: faild to build Int DC") (Maybe DataCon -> DataCon) -> Maybe DataCon -> DataCon
forall a b. (a -> b) -> a -> b
$ do
              (TyConApp iTcNm _) <- TypeView -> Maybe TypeView
forall a. a -> Maybe a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Type -> TypeView
tyView Type
iTy)
              iTc <- UniqMap.lookup iTcNm tcm
              Maybe.listToMaybe (tyConDataCons iTc)

      -- Get elements from vector
          (uniqs1,(vars,elems)) = second (second sconcat . NE.unzip)
                                $ extractElems
                                    uniqs0
                                    is0
                                    consCon
                                    aTy
                                    'I'
                                    n
                                    v

      -- Build a right-biased tree of case-expressions
      let indexed = ((Term, Integer) -> Term -> Term)
-> Term -> [(Term, Integer)] -> Term
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (TyConMap -> DataCon -> Type -> (Term, Integer) -> Term -> Term
index_intElement TyConMap
tcm DataCon
iDc Type
iTy)
                              (Term -> Type -> Term
TyApp (PrimInfo -> Term
Prim PrimInfo
NP.undefined) Type
aTy)
                              ([Term] -> [Integer] -> [(Term, Integer)]
forall a b. [a] -> [b] -> [(a, b)]
zip (NonEmpty Term -> [Term]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty Term
vars) [Integer
0..])
          lb      = [(Id, Term)] -> Term -> Term
Letrec (NonEmpty (Id, Term) -> [(Id, Term)]
forall a. NonEmpty a -> [a]
NE.init NonEmpty (Id, Term)
elems) Term
indexed
      uniqSupply Lens..= uniqs1
      changed lb
  go TyConMap
_ Type
ty = [Char] -> RewriteMonad extra Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> RewriteMonad extra Term)
-> [Char] -> RewriteMonad extra Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"indexReplace_int: argument does not have "
                              [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"a vector type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Type -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr Type
ty

-- | Replace an application of the @Clash.Sized.Vector.dtfold@ primitive on
-- vectors of a known length @n@, by the fully unrolled recursive "definition"
-- of @Clash.Sized.Vector.dtfold@
reduceDTFold
  :: Integer  -- ^ Length of the vector
  -> Type     -- ^ Element type of the argument vector
  -> Term     -- ^ Length of the vector (as a KnownNat)
  -> Term     -- ^ The motive
  -> Term     -- ^ Function to convert elements with
  -> Term     -- ^ Function to combine branches with
  -> Term     -- ^ The vector to fold
  -> TransformContext
  -> NormalizeSession Term
reduceDTFold :: Integer
-> Type
-> Term
-> Term
-> Term
-> Term
-> Term
-> TransformContext
-> NormalizeSession Term
reduceDTFold Integer
n Type
aTy Term
_kn Term
_motive Term
lrFun Term
brFun Term
arg (TransformContext InScopeSet
inScope Context
_ctx) = do
    tcm <- Getting TyConMap RewriteEnv TyConMap
-> RewriteMonad NormalizeState TyConMap
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting TyConMap RewriteEnv TyConMap
Getter RewriteEnv TyConMap
tcCache
    let ty = TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
arg
    go tcm ty
  where
    go :: TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
tcm -> Just Type
ty') = TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
ty'
    go TyConMap
tcm (Type -> TypeView
tyView -> TyConApp TyConName
vecTcNm [Type]
_)
      | (Just TyCon
vecTc) <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
vecTcNm TyConMap
tcm
      , TyConName -> Text
forall a. Name a -> Text
nameOcc TyConName
vecTcNm Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"Clash.Sized.Vector.Vec"
      , [DataCon
_,DataCon
consCon]  <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      = do uniqs0 <- Getting Supply (RewriteState extra) Supply
-> RewriteMonad extra Supply
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting Supply (RewriteState extra) Supply
forall extra (f :: Type -> Type).
Functor f =>
(Supply -> f Supply)
-> RewriteState extra -> f (RewriteState extra)
uniqSupply
           let (uniqs1,(vars,elems)) = second (second sconcat . NE.unzip)
                                     $ extractElems uniqs0 inScope consCon aTy
                                         'T' (2^n) arg
               snatDc = DataCon -> Maybe DataCon -> DataCon
forall a. a -> Maybe a -> a
Maybe.fromMaybe ([Char] -> DataCon
forall a. HasCallStack => [Char] -> a
error [Char]
"reduceDTFold: faild to build SNat") (Maybe DataCon -> DataCon) -> Maybe DataCon -> DataCon
forall a b. (a -> b) -> a -> b
$ do
                  (_ltv:Right snTy:_,_) <- ([Either TyVar Type], Type) -> Maybe ([Either TyVar Type], Type)
forall a. a -> Maybe a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Type -> ([Either TyVar Type], Type)
splitFunForallTy (TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
brFun))
                  (TyConApp snatTcNm _) <- pure (tyView snTy)
                  snatTc <- UniqMap.lookup snatTcNm tcm
                  Maybe.listToMaybe (tyConDataCons snatTc)
               lbody = (Integer -> Term) -> Integer -> [Term] -> Term
doFold (DataCon -> Integer -> Term
buildSNat DataCon
snatDc) (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1) (NonEmpty Term -> [Term]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty Term
vars)
               lb = [(Id, Term)] -> Term -> Term
Letrec (NonEmpty (Id, Term) -> [(Id, Term)]
forall a. NonEmpty a -> [a]
NE.init NonEmpty (Id, Term)
elems) Term
lbody
           uniqSupply Lens..= uniqs1
           changed lb
    go TyConMap
_ Type
ty = [Char] -> RewriteMonad extra Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> RewriteMonad extra Term)
-> [Char] -> RewriteMonad extra Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"reduceDTFold: argument does not have a vector type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Type -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr Type
ty

    doFold :: (Integer -> Term) -> Integer -> [Term] -> Term
    doFold :: (Integer -> Term) -> Integer -> [Term] -> Term
doFold Integer -> Term
_    Integer
_ [Term
x] = Term -> [Either Term Type] -> Term
mkApps Term
lrFun [Term -> Either Term Type
forall a b. a -> Either a b
Left Term
x]
    doFold Integer -> Term
snDc Integer
k [Term]
xs  =
      let ([Term]
xsL,[Term]
xsR) = Int -> [Term] -> ([Term], [Term])
forall a. Int -> [a] -> ([a], [a])
splitAt (Int
2Int -> Integer -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^Integer
k) [Term]
xs
          k' :: Integer
k'        = Integer
kInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1
          eL :: Term
eL        = (Integer -> Term) -> Integer -> [Term] -> Term
doFold Integer -> Term
snDc Integer
k' [Term]
xsL
          eR :: Term
eR        = (Integer -> Term) -> Integer -> [Term] -> Term
doFold Integer -> Term
snDc Integer
k' [Term]
xsR
      in  Term -> [Either Term Type] -> Term
mkApps Term
brFun [Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
k))
                       ,Term -> Either Term Type
forall a b. a -> Either a b
Left  (Integer -> Term
snDc Integer
k)
                       ,Term -> Either Term Type
forall a b. a -> Either a b
Left  Term
eL
                       ,Term -> Either Term Type
forall a b. a -> Either a b
Left  Term
eR
                       ]

-- | Replace an application of the @Clash.Sized.RTree.tdfold@ primitive on
-- trees of a known depth @n@, by the fully unrolled recursive "definition"
-- of @Clash.Sized.RTree.tdfold@
reduceTFold
  :: Integer -- ^ Depth of the tree
  -> Type    -- ^ Element type of the argument tree
  -> Term    -- ^ Depth of the tree (as a KnownNat)
  -> Term    -- ^ The motive
  -> Term    -- ^ Function to convert elements with
  -> Term    -- ^ Function to combine branches with
  -> Term    -- ^ The tree to fold
  -> TransformContext
  -> NormalizeSession Term
reduceTFold :: Integer
-> Type
-> Term
-> Term
-> Term
-> Term
-> Term
-> TransformContext
-> NormalizeSession Term
reduceTFold Integer
n Type
aTy Term
_kn Term
_motive Term
lrFun Term
brFun Term
arg (TransformContext InScopeSet
inScope Context
_ctx) = do
    tcm <- Getting TyConMap RewriteEnv TyConMap
-> RewriteMonad NormalizeState TyConMap
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting TyConMap RewriteEnv TyConMap
Getter RewriteEnv TyConMap
tcCache
    let ty = TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
arg
    go tcm ty
  where
    go :: TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
tcm -> Just Type
ty') = TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
ty'
    go TyConMap
tcm (Type -> TypeView
tyView -> TyConApp TyConName
treeTcNm [Type]
_)
      | (Just TyCon
treeTc) <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
treeTcNm TyConMap
tcm
      , TyConName -> Text
forall a. Name a -> Text
nameOcc TyConName
treeTcNm Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"Clash.Sized.RTree.RTree"
      , [DataCon
lrCon,DataCon
brCon] <- TyCon -> [DataCon]
tyConDataCons TyCon
treeTc
      = do uniqs0 <- Getting Supply (RewriteState extra) Supply
-> RewriteMonad extra Supply
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting Supply (RewriteState extra) Supply
forall extra (f :: Type -> Type).
Functor f =>
(Supply -> f Supply)
-> RewriteState extra -> f (RewriteState extra)
uniqSupply
           let (uniqs1,(vars,elems)) = extractTElems uniqs0 inScope lrCon brCon aTy 'T' n arg
               snatDc = DataCon -> Maybe DataCon -> DataCon
forall a. a -> Maybe a -> a
Maybe.fromMaybe ([Char] -> DataCon
forall a. HasCallStack => [Char] -> a
error [Char]
"reduceTFold: faild to build SNat") (Maybe DataCon -> DataCon) -> Maybe DataCon -> DataCon
forall a b. (a -> b) -> a -> b
$ do
                  (_ltv:Right snTy:_,_) <- ([Either TyVar Type], Type) -> Maybe ([Either TyVar Type], Type)
forall a. a -> Maybe a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Type -> ([Either TyVar Type], Type)
splitFunForallTy (TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
brFun))
                  (TyConApp snatTcNm _) <- pure (tyView snTy)
                  snatTc <- UniqMap.lookup snatTcNm tcm
                  Maybe.listToMaybe (tyConDataCons snatTc)
               lbody = (Integer -> Term) -> Integer -> [Term] -> Term
doFold (DataCon -> Integer -> Term
buildSNat DataCon
snatDc) (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1) [Term]
vars
               lb = ([(Id, Term)] -> Term -> Term
Letrec [(Id, Term)]
elems Term
lbody)
           uniqSupply Lens..= uniqs1
           changed lb
    go TyConMap
_ Type
ty = [Char] -> RewriteMonad extra Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> RewriteMonad extra Term)
-> [Char] -> RewriteMonad extra Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"reduceTFold: argument does not have a tree type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Type -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr Type
ty

    doFold :: (Integer -> Term) -> Integer -> [Term] -> Term
doFold Integer -> Term
_    Integer
_ [Term
x] = Term -> [Either Term Type] -> Term
mkApps Term
lrFun [Term -> Either Term Type
forall a b. a -> Either a b
Left Term
x]
    doFold Integer -> Term
snDc Integer
k [Term]
xs  =
      let ([Term]
xsL,[Term]
xsR) = Int -> [Term] -> ([Term], [Term])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Term] -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [Term]
xs Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) [Term]
xs
          k' :: Integer
k'        = Integer
kInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1
          eL :: Term
eL        = (Integer -> Term) -> Integer -> [Term] -> Term
doFold Integer -> Term
snDc Integer
k' [Term]
xsL
          eR :: Term
eR        = (Integer -> Term) -> Integer -> [Term] -> Term
doFold Integer -> Term
snDc Integer
k' [Term]
xsR
      in  Term -> [Either Term Type] -> Term
mkApps Term
brFun [Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
k))
                       ,Term -> Either Term Type
forall a b. a -> Either a b
Left (Integer -> Term
snDc Integer
k)
                       ,Term -> Either Term Type
forall a b. a -> Either a b
Left Term
eL
                       ,Term -> Either Term Type
forall a b. a -> Either a b
Left Term
eR
                       ]

reduceTReplicate :: Integer -- ^ Depth of the tree
                 -> Type    -- ^ Element type
                 -> Type    -- ^ Result type
                 -> Term    -- ^ Depth of the tree (as an SNat)
                 -> Term    -- ^ Element
                 -> TransformContext
                 -> NormalizeSession Term
reduceTReplicate :: Integer
-> Type
-> Type
-> Term
-> Term
-> TransformContext
-> NormalizeSession Term
reduceTReplicate Integer
n Type
aTy Type
eTy Term
_sn Term
arg TransformContext
_ctx = do
    tcm <- Getting TyConMap RewriteEnv TyConMap
-> RewriteMonad NormalizeState TyConMap
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting TyConMap RewriteEnv TyConMap
Getter RewriteEnv TyConMap
tcCache
    go tcm eTy
  where
    go :: TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
tcm -> Just Type
ty') = TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
ty'
    go TyConMap
tcm (Type -> TypeView
tyView -> TyConApp TyConName
treeTcNm [Type]
_)
      | (Just TyCon
treeTc) <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
treeTcNm TyConMap
tcm
      , TyConName -> Text
forall a. Name a -> Text
nameOcc TyConName
treeTcNm Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"Clash.Sized.RTree.RTree"
      , [DataCon
lrCon,DataCon
brCon] <- TyCon -> [DataCon]
tyConDataCons TyCon
treeTc
      = let retVec :: Term
retVec = DataCon -> DataCon -> Type -> Integer -> [Term] -> Term
mkRTree DataCon
lrCon DataCon
brCon Type
aTy Integer
n (Int -> Term -> [Term]
forall a. Int -> a -> [a]
replicate (Int
2Int -> Integer -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^Integer
n) Term
arg)
        in  Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed Term
retVec
    go TyConMap
_ Type
ty = [Char] -> RewriteMonad extra Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> RewriteMonad extra Term)
-> [Char] -> RewriteMonad extra Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"reduceTReplicate: argument does not have a RTree type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Type -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr Type
ty

buildSNat :: DataCon -> Integer -> Term
buildSNat :: DataCon -> Integer -> Term
buildSNat DataCon
snatDc Integer
i =
  Term -> [Either Term Type] -> Term
mkApps (DataCon -> Term
Data DataCon
snatDc)
         [Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
i))
         ,Term -> Either Term Type
forall a b. a -> Either a b
Left (Literal -> Term
Literal (Integer -> Literal
NaturalLiteral (Integer -> Integer
forall a. Integral a => a -> Integer
toInteger Integer
i)))
         ]