{-|
  Copyright  :  (C) 2015-2016, University of Twente,
                    2016     , Myrtle Software Ltd
  License    :  BSD2 (see the file LICENSE)
  Maintainer :  Christiaan Baaij <christiaan.baaij@gmail.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           Data.List                        (mapAccumR)
import           Data.List.Extra                  (zipEqual)
import qualified Data.Maybe                       as Maybe
import           TextShow                         (showt)

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

import           Clash.Core.DataCon               (DataCon)
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)
import           Clash.Core.TermInfo
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, undefinedTm)
import           Clash.Core.Var                   (Var (..), mkTyVar, mkLocalId)
import           Clash.Core.VarEnv
  (InScopeSet, extendInScopeSetList)
import {-# SOURCE #-} Clash.Normalize.Strategy
import           Clash.Normalize.Types
import           Clash.Rewrite.Types
import           Clash.Rewrite.Util
import           Clash.Unique
import           Clash.Util
import qualified Clash.Util.Interpolate           as I

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

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

typeNatSub :: TyConName
typeNatSub :: TyConName
typeNatSub =
  NameSort -> OccName -> Unique -> SrcSpan -> TyConName
forall a. NameSort -> OccName -> Unique -> SrcSpan -> Name a
Name NameSort
User OccName
"GHC.TypeNats.-" (Unique -> Unique
getKey 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 (OccName -> Type -> WorkInfo -> IsMultiPrim -> PrimInfo
PrimInfo OccName
"Clash.Sized.Vector.head" (TyConName -> Type
vecHeadTy TyConName
vecTcNm) WorkInfo
WorkNever IsMultiPrim
SingleResult)

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 (OccName -> Type -> WorkInfo -> IsMultiPrim -> PrimInfo
PrimInfo OccName
"Clash.Sized.Vector.last" (TyConName -> Type
vecHeadTy TyConName
vecTcNm) WorkInfo
WorkNever IsMultiPrim
SingleResult)

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 (OccName -> Unique -> TyName
forall a. OccName -> Unique -> Name a
mkUnsafeSystemName OccName
"a" Unique
0)
  nTV :: TyVar
nTV = Type -> TyName -> TyVar
mkTyVar Type
typeNatKind (OccName -> Unique -> TyName
forall a. OccName -> Unique -> Name a
mkUnsafeSystemName OccName
"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 (OccName -> Type -> WorkInfo -> IsMultiPrim -> PrimInfo
PrimInfo OccName
"Clash.Sized.Vector.tail" (TyConName -> Type
vecTailTy TyConName
vecTcNm) WorkInfo
WorkNever IsMultiPrim
SingleResult)

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 (OccName -> Type -> WorkInfo -> IsMultiPrim -> PrimInfo
PrimInfo OccName
"Clash.Sized.Vector.init" (TyConName -> Type
vecTailTy TyConName
vecTcNm) WorkInfo
WorkNever IsMultiPrim
SingleResult)

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 (OccName -> Unique -> TyName
forall a. OccName -> Unique -> Name a
mkUnsafeSystemName OccName
"n" Unique
0)
  aTV :: TyVar
aTV = Type -> TyName -> TyVar
mkTyVar Type
liftedTypeKind (OccName -> Unique -> TyName
forall a. OccName -> Unique -> Name a
mkUnsafeSystemName OccName
"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
  -> 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 =
  ( 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)] )
 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)))]
  Just [Type
coTy, Type
_elTy, Type
restTy] = DataCon -> [Type] -> Maybe [Type]
dataConInstArgTys DataCon
consCon [Type]
tys

  mTV :: TyVar
mTV = Type -> TyName -> TyVar
mkTyVar Type
typeNatKind (OccName -> Unique -> TyName
forall a. OccName -> Unique -> Name a
mkUnsafeSystemName OccName
"m" Unique
0)
  co :: Id
co = Type -> TmName -> Id
mkLocalId Type
coTy (OccName -> Unique -> TmName
forall a. OccName -> Unique -> Name a
mkUnsafeSystemName OccName
"_co_" Unique
1)
  el :: Id
el = Type -> TmName -> Id
mkLocalId Type
elTy (OccName -> Unique -> TmName
forall a. OccName -> Unique -> Name a
mkUnsafeSystemName OccName
"el" Unique
2)
  rest :: Id
rest = Type -> TmName -> Id
mkLocalId Type
restTy (OccName -> Unique -> TmName
forall a. OccName -> Unique -> Name a
mkUnsafeSystemName OccName
"res" Unique
3)

  pat :: Pat
pat = DataCon -> [TyVar] -> [Id] -> Pat
DataPat DataCon
consCon [TyVar
mTV] [Id
co, Id
el, Id
rest]

-- Make case statement that projects the _head_ from a given vector
extractHead
  :: DataCon
  -- ^ The Cons (:>) constructor
  -> Type
  -- ^ Element type
  -> Integer
  -- ^ Length of the vector
  -> Term
  -- ^ Vector to extract head from
  -> Term
  -- ^ Head of vector
extractHead :: DataCon -> Type -> Integer -> Term -> Term
extractHead DataCon
consCon Type
elTy Integer
vLength Term
vec =
  (Term, Term) -> Term
forall a b. (a, b) -> a
fst (DataCon -> Type -> Integer -> Term -> (Term, Term)
extractHeadTail DataCon
consCon Type
elTy Integer
vLength Term
vec)

-- Make case statement that projects the _tail_ from a given vector
extractTail
  :: DataCon
  -- ^ The Cons (:>) constructor
  -> Type
  -- ^ Element type
  -> Integer
  -- ^ Length of the vector
  -> Term
  -- ^ Vector to extract head from
  -> Term
  -- ^ Tail of vector
extractTail :: DataCon -> Type -> Integer -> Term -> Term
extractTail DataCon
consCon Type
elTy Integer
vLength Term
vec =
  (Term, Term) -> Term
forall a b. (a, b) -> b
snd (DataCon -> Type -> Integer -> Term -> (Term, Term)
extractHeadTail DataCon
consCon Type
elTy Integer
vLength Term
vec)

-- | 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 :: 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 =
    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 ]

 where
  args :: Maybe [Type]
args = 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))]
  Just (Type
consCoTy : [Type]
_) = Maybe [Type]
args

-- | Create an empty vector
mkVecNil
  :: DataCon
  -- ^ The Nil constructor
  -> Type
  -- ^ The element type
  -> Term
mkVecNil :: DataCon -> Type -> Term
mkVecNil DataCon
nilCon Type
resTy =
  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) ]
 where
  args :: Maybe [Type]
args = DataCon -> [Type] -> Maybe [Type]
dataConInstArgTys DataCon
nilCon [LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
0), Type
resTy]
  Just (Type
nilCoTy : [Type]
_ ) = Maybe [Type]
args

-- | 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
  :: InScopeSet
  -> Integer
  -- ^ Length of the vector
  -> Type
  -- ^ Element of type of the vector
  -> Term
  -- ^ The vector to reverse
  -> NormalizeSession Term
reduceReverse :: InScopeSet -> Integer -> Type -> Term -> NormalizeSession Term
reduceReverse InScopeSet
inScope0 Integer
n Type
elTy Term
vArg = do
  TyConMap
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
Lens' RewriteEnv TyConMap
tcCache
  let ty :: Type
ty = TyConMap -> Term -> Type
termType TyConMap
tcm Term
vArg
  TyConMap -> Type -> NormalizeSession Term
forall extra. TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
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
lookupUniqMap TyConName
vecTcNm TyConMap
tcm
      , TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
vecTcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"Clash.Sized.Vector.Vec"
    , [DataCon
nilCon, DataCon
consCon] <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
    = do
      Supply
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. Lens' (RewriteState extra) Supply
uniqSupply
      let (Supply
uniqs1,([Term]
vars,[LetBinding]
elems)) = ([(Term, [LetBinding])] -> ([Term], [LetBinding]))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (([[LetBinding]] -> [LetBinding])
-> ([Term], [[LetBinding]]) -> ([Term], [LetBinding])
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second [[LetBinding]] -> [LetBinding]
forall (t :: Type -> Type) a. Foldable t => t [a] -> [a]
concat (([Term], [[LetBinding]]) -> ([Term], [LetBinding]))
-> ([(Term, [LetBinding])] -> ([Term], [[LetBinding]]))
-> [(Term, [LetBinding])]
-> ([Term], [LetBinding])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Term, [LetBinding])] -> ([Term], [[LetBinding]])
forall a b. [(a, b)] -> ([a], [b])
unzip)
                                ((Supply, [(Term, [LetBinding])])
 -> (Supply, ([Term], [LetBinding])))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall a b. (a -> b) -> a -> b
$ Supply
-> InScopeSet
-> DataCon
-> Type
-> Char
-> Integer
-> Term
-> (Supply, [(Term, [LetBinding])])
extractElems Supply
uniqs0 InScopeSet
inScope0 DataCon
consCon Type
elTy Char
'V' Integer
n Term
vArg
          lbody :: Term
lbody = DataCon -> DataCon -> Type -> Integer -> [Term] -> Term
mkVec DataCon
nilCon DataCon
consCon Type
elTy Integer
n ([Term] -> [Term]
forall a. [a] -> [a]
reverse [Term]
vars)
          lb :: Term
lb    = [LetBinding] -> Term -> Term
Letrec ([LetBinding] -> [LetBinding]
forall a. [a] -> [a]
init [LetBinding]
elems) Term
lbody
      (Supply -> Identity Supply)
-> RewriteState extra -> Identity (RewriteState extra)
forall extra. Lens' (RewriteState extra) Supply
uniqSupply ((Supply -> Identity Supply)
 -> RewriteState extra -> Identity (RewriteState extra))
-> Supply -> RewriteMonad extra ()
forall s (m :: Type -> Type) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
Lens..= Supply
uniqs1
      Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed Term
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
$ $([Char]
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
  :: TransformContext
  -> 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
  -> NormalizeSession Term
reduceZipWith :: TransformContext
-> PrimInfo
-> Integer
-> Type
-> Type
-> Type
-> Term
-> Term
-> Term
-> NormalizeSession Term
reduceZipWith TransformContext
_ctx PrimInfo
zipWithPrimInfo Integer
n Type
lhsElTy Type
rhsElTy Type
resElTy Term
fun Term
lhsArg Term
rhsArg = do
  TyConMap
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
Lens' RewriteEnv TyConMap
tcCache
  Term -> NormalizeSession Term
forall a extra. a -> RewriteMonad extra a
changed (TyConMap -> Type -> Term
go TyConMap
tcm (TyConMap -> Term -> Type
termType TyConMap
tcm Term
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
lookupUniqMap TyConName
vecTcNm TyConMap
tcm
    , TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
vecTcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"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
$ $([Char]
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
  :: TransformContext
  -> 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
  -> NormalizeSession Term
reduceMap :: TransformContext
-> PrimInfo
-> Integer
-> Type
-> Type
-> Term
-> Term
-> NormalizeSession Term
reduceMap TransformContext
_ctx PrimInfo
mapPrimInfo Integer
n Type
argElTy Type
resElTy Term
fun Term
arg = do
    TyConMap
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
Lens' RewriteEnv TyConMap
tcCache
    let ty :: Type
ty = TyConMap -> Term -> Type
termType TyConMap
tcm Term
arg
    Term -> NormalizeSession Term
forall a extra. a -> RewriteMonad extra a
changed (TyConMap -> Type -> Term
go TyConMap
tcm Type
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
lookupUniqMap TyConName
vecTcNm TyConMap
tcm
      , TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
vecTcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"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
$ $([Char]
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
  :: TransformContext
  -> Integer  -- ^ Length of the vector
  -> Type -- ^ Argument type of the function
  -> Type -- ^ Result type of the function
  -> Term -- ^ The imap'd function
  -> Term -- ^ The imap'd over vector
  -> NormalizeSession Term
reduceImap :: TransformContext
-> Integer -> Type -> Type -> Term -> Term -> NormalizeSession Term
reduceImap (TransformContext InScopeSet
is0 Context
ctx) Integer
n Type
argElTy Type
resElTy Term
fun Term
arg = do
    TyConMap
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
Lens' RewriteEnv TyConMap
tcCache
    let ty :: Type
ty = TyConMap -> Term -> Type
termType TyConMap
tcm Term
arg
    TyConMap -> Type -> NormalizeSession Term
go TyConMap
tcm Type
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
lookupUniqMap TyConName
vecTcNm TyConMap
tcm
      , TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
vecTcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"Clash.Sized.Vector.Vec"
      , [DataCon
nilCon,DataCon
consCon] <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      = do
        Supply
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. Lens' (RewriteState extra) Supply
uniqSupply
        Term
fun1 <- NormRewrite
constantPropagation (InScopeSet -> Context -> TransformContext
TransformContext InScopeSet
is0 (Maybe (OccName, Unique, Unique) -> CoreContext
AppArg Maybe (OccName, Unique, Unique)
forall a. Maybe a
NothingCoreContext -> Context -> Context
forall a. a -> [a] -> [a]
:Context
ctx)) Term
fun
        let is1 :: InScopeSet
is1 = InScopeSet -> [Id] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 (Term -> [Id]
collectTermIds Term
fun1)
            ((Supply, InScopeSet)
uniqs1,TyVar
nTv) = (Supply, InScopeSet)
-> (OccName, Type) -> ((Supply, InScopeSet), TyVar)
mkUniqSystemTyVar (Supply
uniqs0,InScopeSet
is1) (OccName
"n",Type
typeNatKind)
            (Supply
uniqs2,([Term]
vars,[LetBinding]
elems)) = ([(Term, [LetBinding])] -> ([Term], [LetBinding]))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (([[LetBinding]] -> [LetBinding])
-> ([Term], [[LetBinding]]) -> ([Term], [LetBinding])
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second [[LetBinding]] -> [LetBinding]
forall (t :: Type -> Type) a. Foldable t => t [a] -> [a]
concat (([Term], [[LetBinding]]) -> ([Term], [LetBinding]))
-> ([(Term, [LetBinding])] -> ([Term], [[LetBinding]]))
-> [(Term, [LetBinding])]
-> ([Term], [LetBinding])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Term, [LetBinding])] -> ([Term], [[LetBinding]])
forall a b. [(a, b)] -> ([a], [b])
unzip)
                                  ((Supply, [(Term, [LetBinding])])
 -> (Supply, ([Term], [LetBinding])))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall a b. (a -> b) -> a -> b
$ (Supply
 -> InScopeSet
 -> DataCon
 -> Type
 -> Char
 -> Integer
 -> Term
 -> (Supply, [(Term, [LetBinding])]))
-> (Supply, InScopeSet)
-> DataCon
-> Type
-> Char
-> Integer
-> Term
-> (Supply, [(Term, [LetBinding])])
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Supply
-> InScopeSet
-> DataCon
-> Type
-> Char
-> Integer
-> Term
-> (Supply, [(Term, [LetBinding])])
extractElems (Supply, InScopeSet)
uniqs1 DataCon
consCon Type
argElTy Char
'I' Integer
n Term
arg
            (Right Type
idxTy:[Either TyVar Type]
_,Type
_) = Type -> ([Either TyVar Type], Type)
splitFunForallTy (TyConMap -> Term -> Type
termType TyConMap
tcm Term
fun)
            (TyConApp TyConName
idxTcNm [Type]
_) = Type -> TypeView
tyView Type
idxTy
            -- fromInteger# :: KnownNat n => Integer -> Index n
            idxFromIntegerTy :: Type
idxFromIntegerTy = TyVar -> Type -> Type
ForAllTy TyVar
nTv
                                        ((Type -> Type -> Type) -> Type -> [Type] -> Type
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 :: Term
idxFromInteger   = PrimInfo -> Term
Prim (OccName -> Type -> WorkInfo -> IsMultiPrim -> PrimInfo
PrimInfo OccName
"Clash.Sized.Internal.Index.fromInteger#" Type
idxFromIntegerTy WorkInfo
WorkNever IsMultiPrim
SingleResult)
            idxs :: [Term]
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]
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 [Term]
vars
            lbody :: Term
lbody            = DataCon -> DataCon -> Type -> Integer -> [Term] -> Term
mkVec DataCon
nilCon DataCon
consCon Type
resElTy Integer
n [Term]
funApps
            lb :: Term
lb               = [LetBinding] -> Term -> Term
Letrec ([LetBinding] -> [LetBinding]
forall a. [a] -> [a]
init [LetBinding]
elems) Term
lbody
        (Supply -> Identity Supply)
-> RewriteState NormalizeState
-> Identity (RewriteState NormalizeState)
forall extra. Lens' (RewriteState extra) Supply
uniqSupply ((Supply -> Identity Supply)
 -> RewriteState NormalizeState
 -> Identity (RewriteState NormalizeState))
-> Supply -> RewriteMonad NormalizeState ()
forall s (m :: Type -> Type) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
Lens..= Supply
uniqs2
        Term -> NormalizeSession Term
forall a extra. a -> RewriteMonad extra a
changed Term
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
$ $([Char]
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
  :: TransformContext
  -> Integer
  -- ^ Length of vector
  -> Type
  -- ^ Vector's element type
  -> Type
  -- ^ Vector's type
  -> Term
  -- ^ iterateI's HO-function argument
  -> Term
  -- ^ iterateI's start value
  -> RewriteMonad NormalizeState Term
  -- ^ Fully unrolled definition
reduceIterateI :: TransformContext
-> Integer -> Type -> Type -> Term -> Term -> NormalizeSession Term
reduceIterateI (TransformContext InScopeSet
is0 Context
ctx) Integer
n Type
aTy Type
vTy Term
f0 Term
a = do
  TyConMap
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
Lens' RewriteEnv TyConMap
tcCache
  Term
f1 <- NormRewrite
constantPropagation (InScopeSet -> Context -> TransformContext
TransformContext InScopeSet
is0 (Maybe (OccName, Unique, Unique) -> CoreContext
AppArg Maybe (OccName, Unique, Unique)
forall a. Maybe a
NothingCoreContext -> Context -> Context
forall a. a -> [a] -> [a]
:Context
ctx)) Term
f0

  -- Generate uniq ids for element assignments.
  Supply
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. Lens' (RewriteState extra) Supply
uniqSupply
  let
    is1 :: InScopeSet
is1 = InScopeSet -> [Id] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 (Term -> [Id]
collectTermIds Term
f1)
    ((Supply
uniqs1, InScopeSet
_is2), [Id]
elementIds) =
      ((Supply, InScopeSet)
 -> (OccName, Type) -> ((Supply, InScopeSet), Id))
-> (Supply, InScopeSet)
-> [(OccName, Type)]
-> ((Supply, InScopeSet), [Id])
forall (t :: Type -> Type) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumR
        (Supply, InScopeSet)
-> (OccName, Type) -> ((Supply, InScopeSet), Id)
mkUniqInternalId
        (Supply
uniqs0, InScopeSet
is1)
        ([OccName] -> [Type] -> [(OccName, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Integer -> OccName) -> [Integer] -> [OccName]
forall a b. (a -> b) -> [a] -> [b]
map ((OccName
"el" OccName -> OccName -> OccName
forall a. Semigroup a => a -> a -> a
<>) (OccName -> OccName) -> (Integer -> OccName) -> Integer -> OccName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> OccName
forall a. TextShow a => a -> OccName
showt) [Integer
1..Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1]) (Type -> [Type]
forall a. a -> [a]
repeat Type
aTy))
  (Supply -> Identity Supply)
-> RewriteState NormalizeState
-> Identity (RewriteState NormalizeState)
forall extra. Lens' (RewriteState extra) Supply
uniqSupply ((Supply -> Identity Supply)
 -> RewriteState NormalizeState
 -> Identity (RewriteState NormalizeState))
-> Supply -> RewriteMonad NormalizeState ()
forall s (m :: Type -> Type) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
.= Supply
uniqs1

  let
    TyConApp TyConName
vecTcNm [Type]
_ = Type -> TypeView
tyView Type
vTy
    Just TyCon
vecTc = TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
lookupUniqMap TyConName
vecTcNm TyConMap
tcm
    [DataCon
nilCon, DataCon
consCon] = TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
    elems :: [Term]
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
vec = DataCon -> DataCon -> Type -> Integer -> [Term] -> Term
mkVec DataCon
nilCon DataCon
consCon Type
aTy Integer
n (Unique -> [Term] -> [Term]
forall a. Unique -> [a] -> [a]
take (Integer -> Unique
forall a. Num a => Integer -> a
fromInteger Integer
n) (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))

  -- Result:
  --   let
  --     el1 = f a
  --     el2 = f el1
  --     el3 = f el2
  --     ..
  --   in
  --     (a :> el1 :> el2 :> el3 :> ..)
  --
  Term -> NormalizeSession Term
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ([LetBinding] -> Term -> Term
Letrec ([Id] -> [Term] -> [LetBinding]
forall a b. [a] -> [b] -> [(a, b)]
zip [Id]
elementIds [Term]
elems) Term
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
  :: TransformContext
  -> Integer  -- ^ Length of the vector
  -> 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
  -> NormalizeSession Term
reduceTraverse :: TransformContext
-> Integer
-> Type
-> Type
-> Type
-> Term
-> Term
-> Term
-> NormalizeSession Term
reduceTraverse (TransformContext InScopeSet
is0 Context
ctx) Integer
n Type
aTy Type
fTy Type
bTy Term
dict Term
fun Term
arg = do
    TyConMap
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
Lens' RewriteEnv TyConMap
tcCache
    let (TyConApp TyConName
apDictTcNm [Type]
_) = Type -> TypeView
tyView (TyConMap -> Term -> Type
termType TyConMap
tcm Term
dict)
        ty :: Type
ty = TyConMap -> Term -> Type
termType TyConMap
tcm Term
arg
    TyConMap -> TyConName -> Type -> NormalizeSession Term
forall a.
Uniquable a =>
TyConMap -> a -> Type -> NormalizeSession Term
go TyConMap
tcm TyConName
apDictTcNm Type
ty
  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
lookupUniqMap TyConName
vecTcNm TyConMap
tcm
      , TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
vecTcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"Clash.Sized.Vector.Vec"
      , [DataCon
nilCon,DataCon
consCon] <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      = do
        Supply
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. Lens' (RewriteState extra) Supply
uniqSupply
        Term
fun1 <- NormRewrite
constantPropagation (InScopeSet -> Context -> TransformContext
TransformContext InScopeSet
is0 (Maybe (OccName, Unique, Unique) -> CoreContext
AppArg Maybe (OccName, Unique, Unique)
forall a. Maybe a
NothingCoreContext -> Context -> Context
forall a. a -> [a] -> [a]
:Context
ctx)) Term
fun
        let is1 :: InScopeSet
is1 = InScopeSet -> [Id] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 (Term -> [Id]
collectTermIds Term
fun1)
            (Just TyCon
apDictTc)    = a -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
lookupUniqMap a
apDictTcNm TyConMap
tcm
            [DataCon
apDictCon]        = TyCon -> [DataCon]
tyConDataCons TyCon
apDictTc
            (Just [Type]
apDictIdTys) = DataCon -> [Type] -> Maybe [Type]
dataConInstArgTys DataCon
apDictCon [Type
fTy]
            ((Supply, InScopeSet)
uniqs1,apDictIds :: [Id]
apDictIds@[Id
functorDictId,Id
pureId,Id
apId,Id
_,Id
_,Id
_]) =
              ((Supply, InScopeSet)
 -> (OccName, Type) -> ((Supply, InScopeSet), Id))
-> (Supply, InScopeSet)
-> [(OccName, Type)]
-> ((Supply, InScopeSet), [Id])
forall (t :: Type -> Type) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumR (Supply, InScopeSet)
-> (OccName, Type) -> ((Supply, InScopeSet), Id)
mkUniqInternalId (Supply
uniqs0,InScopeSet
is1)
                ([OccName] -> [Type] -> [(OccName, Type)]
forall a b. HasCallStack => [a] -> [b] -> [(a, b)]
zipEqual [OccName
"functorDict",OccName
"pure",OccName
"ap",OccName
"liftA2",OccName
"apConstL",OccName
"apConstR"]
                     [Type]
apDictIdTys)

            (TyConApp TyConName
funcDictTcNm [Type]
_) = Type -> TypeView
tyView ([Type] -> Type
forall a. [a] -> a
head [Type]
apDictIdTys)
            (Just TyCon
funcDictTc) = TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
lookupUniqMap TyConName
funcDictTcNm TyConMap
tcm
            [DataCon
funcDictCon] = TyCon -> [DataCon]
tyConDataCons TyCon
funcDictTc
            (Just [Type]
funcDictIdTys) = DataCon -> [Type] -> Maybe [Type]
dataConInstArgTys DataCon
funcDictCon [Type
fTy]
            ((Supply, InScopeSet)
uniqs2,funcDicIds :: [Id]
funcDicIds@[Id
fmapId,Id
_]) =
              ((Supply, InScopeSet)
 -> (OccName, Type) -> ((Supply, InScopeSet), Id))
-> (Supply, InScopeSet)
-> [(OccName, Type)]
-> ((Supply, InScopeSet), [Id])
forall (t :: Type -> Type) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumR (Supply, InScopeSet)
-> (OccName, Type) -> ((Supply, InScopeSet), Id)
mkUniqInternalId (Supply, InScopeSet)
uniqs1
                ([OccName] -> [Type] -> [(OccName, Type)]
forall a b. HasCallStack => [a] -> [b] -> [(a, b)]
zipEqual [OccName
"fmap",OccName
"fmapConst"] [Type]
funcDictIdTys)

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

            -- Extract the 'pure' function from the Applicative dictionary
            pureTy :: Type
pureTy = Id -> Type
forall a. Var a -> Type
varType Id
pureId
            pureTm :: Term
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 :: Type
apTy   = Id -> Type
forall a. Var a -> Type
varType Id
apId
            apTm :: Term
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 :: Type
funcTy = Id -> Type
forall a. Var a -> Type
varType Id
functorDictId
            funcTm :: Term
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 :: Type
fmapTy = Id -> Type
forall a. Var a -> Type
varType Id
fmapId
            fmapTm :: Term
fmapTm = Term -> Type -> [Alt] -> Term
Case (Id -> Term
Var Id
functorDictId) Type
fmapTy
                          [(Pat
fnPat, Id -> Term
Var Id
fmapId)]

            (Supply
uniqs3,([Term]
vars,[LetBinding]
elems)) = ([(Term, [LetBinding])] -> ([Term], [LetBinding]))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (([[LetBinding]] -> [LetBinding])
-> ([Term], [[LetBinding]]) -> ([Term], [LetBinding])
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second [[LetBinding]] -> [LetBinding]
forall (t :: Type -> Type) a. Foldable t => t [a] -> [a]
concat (([Term], [[LetBinding]]) -> ([Term], [LetBinding]))
-> ([(Term, [LetBinding])] -> ([Term], [[LetBinding]]))
-> [(Term, [LetBinding])]
-> ([Term], [LetBinding])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Term, [LetBinding])] -> ([Term], [[LetBinding]])
forall a b. [(a, b)] -> ([a], [b])
unzip)
                                  ((Supply, [(Term, [LetBinding])])
 -> (Supply, ([Term], [LetBinding])))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall a b. (a -> b) -> a -> b
$ (Supply
 -> InScopeSet
 -> DataCon
 -> Type
 -> Char
 -> Integer
 -> Term
 -> (Supply, [(Term, [LetBinding])]))
-> (Supply, InScopeSet)
-> DataCon
-> Type
-> Char
-> Integer
-> Term
-> (Supply, [(Term, [LetBinding])])
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Supply
-> InScopeSet
-> DataCon
-> Type
-> Char
-> Integer
-> Term
-> (Supply, [(Term, [LetBinding])])
extractElems (Supply, InScopeSet)
uniqs2 DataCon
consCon Type
aTy Char
'T' Integer
n Term
arg

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

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

            lb :: Term
lb      = [LetBinding] -> Term -> Term
Letrec ([(([Id]
apDictIds[Id] -> Unique -> Id
forall a. [a] -> Unique -> a
!!Unique
0), Term
funcTm)
                              ,(([Id]
apDictIds[Id] -> Unique -> Id
forall a. [a] -> Unique -> a
!!Unique
1), Term
pureTm)
                              ,(([Id]
apDictIds[Id] -> Unique -> Id
forall a. [a] -> Unique -> a
!!Unique
2), Term
apTm)
                              ,(([Id]
funcDicIds[Id] -> Unique -> Id
forall a. [a] -> Unique -> a
!!Unique
0), Term
fmapTm)
                              ] [LetBinding] -> [LetBinding] -> [LetBinding]
forall a. [a] -> [a] -> [a]
++ [LetBinding] -> [LetBinding]
forall a. [a] -> [a]
init [LetBinding]
elems) Term
lbody
        (Supply -> Identity Supply)
-> RewriteState NormalizeState
-> Identity (RewriteState NormalizeState)
forall extra. Lens' (RewriteState extra) Supply
uniqSupply ((Supply -> Identity Supply)
 -> RewriteState NormalizeState
 -> Identity (RewriteState NormalizeState))
-> Supply -> RewriteMonad NormalizeState ()
forall s (m :: Type -> Type) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
Lens..= Supply
uniqs3
        Term -> NormalizeSession Term
forall a extra. a -> RewriteMonad extra a
changed Term
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
$ $([Char]
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 = [Type] -> Type
forall a. [a] -> a
head (Maybe [Type] -> [Type]
forall a. HasCallStack => Maybe a -> a
Maybe.fromJust (DataCon -> [Type] -> Maybe [Type]
dataConInstArgTys DataCon
nilCon [(LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
0))
                                                             ,Type
bTy]))

    consCoTy :: Integer -> Type
consCoTy Integer
n = [Type] -> Type
forall a. [a] -> a
head (Maybe [Type] -> [Type]
forall a. HasCallStack => Maybe a -> a
Maybe.fromJust (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)))]))

-- | 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
  :: TransformContext
  -> 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
  -> NormalizeSession Term
reduceFoldr :: TransformContext
-> PrimInfo
-> Integer
-> Type
-> Term
-> Term
-> Term
-> NormalizeSession Term
reduceFoldr TransformContext
_ PrimInfo
_ Integer
0 Type
_ Term
_ Term
start Term
_ = Term -> NormalizeSession Term
forall a extra. a -> RewriteMonad extra a
changed Term
start
reduceFoldr TransformContext
_ctx PrimInfo
foldrPrimInfo Integer
n Type
aTy Term
fun Term
start Term
arg = do
    TyConMap
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
Lens' RewriteEnv TyConMap
tcCache
    let ty :: Type
ty = TyConMap -> Term -> Type
termType TyConMap
tcm Term
arg
    Term -> NormalizeSession Term
forall a extra. a -> RewriteMonad extra a
changed (TyConMap -> Type -> Term
go TyConMap
tcm Type
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 -> OccName
forall a. Name a -> OccName
nameOcc TyConName
vecTcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"Clash.Sized.Vector.Vec"
      , Just TyCon
vecTc <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
lookupUniqMap 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
termType 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
$ $([Char]
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
  :: TransformContext
  -> Integer
  -- ^ Length of the vector
  -> Type
  -- ^ Element type of the argument vector
  -> Term
  -- ^ The function to fold with
  -> Term
  -- ^ The argument vector
  -> NormalizeSession Term
reduceFold :: TransformContext
-> Integer -> Type -> Term -> Term -> NormalizeSession Term
reduceFold (TransformContext InScopeSet
is0 Context
ctx) Integer
n Type
aTy Term
fun Term
arg = do
    TyConMap
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
Lens' RewriteEnv TyConMap
tcCache
    let ty :: Type
ty = TyConMap -> Term -> Type
termType TyConMap
tcm Term
arg
    TyConMap -> Type -> NormalizeSession Term
go TyConMap
tcm Type
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
lookupUniqMap TyConName
vecTcNm TyConMap
tcm
      , TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
vecTcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"Clash.Sized.Vector.Vec"
      , [DataCon
_,DataCon
consCon]  <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      = do
        Supply
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. Lens' (RewriteState extra) Supply
uniqSupply
        Term
fun1 <- NormRewrite
constantPropagation (InScopeSet -> Context -> TransformContext
TransformContext InScopeSet
is0 (Maybe (OccName, Unique, Unique) -> CoreContext
AppArg Maybe (OccName, Unique, Unique)
forall a. Maybe a
NothingCoreContext -> Context -> Context
forall a. a -> [a] -> [a]
:Context
ctx)) Term
fun
        let is1 :: InScopeSet
is1 = InScopeSet -> [Id] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 (Term -> [Id]
collectTermIds Term
fun1)
            (Supply
uniqs1,([Term]
vars,[LetBinding]
elems)) = ([(Term, [LetBinding])] -> ([Term], [LetBinding]))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (([[LetBinding]] -> [LetBinding])
-> ([Term], [[LetBinding]]) -> ([Term], [LetBinding])
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second [[LetBinding]] -> [LetBinding]
forall (t :: Type -> Type) a. Foldable t => t [a] -> [a]
concat (([Term], [[LetBinding]]) -> ([Term], [LetBinding]))
-> ([(Term, [LetBinding])] -> ([Term], [[LetBinding]]))
-> [(Term, [LetBinding])]
-> ([Term], [LetBinding])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Term, [LetBinding])] -> ([Term], [[LetBinding]])
forall a b. [(a, b)] -> ([a], [b])
unzip)
                                  ((Supply, [(Term, [LetBinding])])
 -> (Supply, ([Term], [LetBinding])))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall a b. (a -> b) -> a -> b
$ Supply
-> InScopeSet
-> DataCon
-> Type
-> Char
-> Integer
-> Term
-> (Supply, [(Term, [LetBinding])])
extractElems Supply
uniqs0 InScopeSet
is1 DataCon
consCon Type
aTy Char
'F' Integer
n Term
arg
            lbody :: Term
lbody            = Term -> [Term] -> Term
foldV Term
fun1 [Term]
vars
            lb :: Term
lb               = [LetBinding] -> Term -> Term
Letrec ([LetBinding] -> [LetBinding]
forall a. [a] -> [a]
init [LetBinding]
elems) Term
lbody
        (Supply -> Identity Supply)
-> RewriteState NormalizeState
-> Identity (RewriteState NormalizeState)
forall extra. Lens' (RewriteState extra) Supply
uniqSupply ((Supply -> Identity Supply)
 -> RewriteState NormalizeState
 -> Identity (RewriteState NormalizeState))
-> Supply -> RewriteMonad NormalizeState ()
forall s (m :: Type -> Type) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
Lens..= Supply
uniqs1
        Term -> NormalizeSession Term
forall a extra. a -> RewriteMonad extra a
changed Term
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
$ $([Char]
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) = Unique -> [Term] -> ([Term], [Term])
forall a. Unique -> [a] -> ([a], [a])
splitAt ([Term] -> Unique
forall (t :: Type -> Type) a. Foldable t => t a -> Unique
length [Term]
as Unique -> Unique -> Unique
forall a. Integral a => a -> a -> a
`div` Unique
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
  :: InScopeSet
  -> Integer
  -- ^ Length of the vector
  -> Type
  -- ^ Element type of the argument vector
  -> Term
  -- ^ Function to fold with
  -> Term
  -- ^ Starting value
  -> Term
  -- ^ The vector to fold
  -> NormalizeSession Term
reduceDFold :: InScopeSet
-> Integer -> Type -> Term -> Term -> Term -> NormalizeSession Term
reduceDFold InScopeSet
_ Integer
0 Type
_ Term
_ Term
start Term
_ = Term -> NormalizeSession Term
forall a extra. a -> RewriteMonad extra a
changed Term
start
reduceDFold InScopeSet
is0 Integer
n Type
aTy Term
fun Term
start Term
arg = do
    TyConMap
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
Lens' RewriteEnv TyConMap
tcCache
    let ty :: Type
ty = TyConMap -> Term -> Type
termType TyConMap
tcm Term
arg
    TyConMap -> Type -> NormalizeSession Term
forall extra. TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
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
lookupUniqMap TyConName
vecTcNm TyConMap
tcm
      , TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
vecTcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"Clash.Sized.Vector.Vec"
      , [DataCon
_,DataCon
consCon]  <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      = do
        Supply
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. Lens' (RewriteState extra) Supply
uniqSupply
        let is1 :: InScopeSet
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.
            (Supply
uniqs1,([Term]
vars,[LetBinding]
elems)) = ([(Term, [LetBinding])] -> ([Term], [LetBinding]))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (([[LetBinding]] -> [LetBinding])
-> ([Term], [[LetBinding]]) -> ([Term], [LetBinding])
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second [[LetBinding]] -> [LetBinding]
forall (t :: Type -> Type) a. Foldable t => t [a] -> [a]
concat (([Term], [[LetBinding]]) -> ([Term], [LetBinding]))
-> ([(Term, [LetBinding])] -> ([Term], [[LetBinding]]))
-> [(Term, [LetBinding])]
-> ([Term], [LetBinding])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Term, [LetBinding])] -> ([Term], [[LetBinding]])
forall a b. [(a, b)] -> ([a], [b])
unzip)
                                  ((Supply, [(Term, [LetBinding])])
 -> (Supply, ([Term], [LetBinding])))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall a b. (a -> b) -> a -> b
$ Supply
-> InScopeSet
-> DataCon
-> Type
-> Char
-> Integer
-> Term
-> (Supply, [(Term, [LetBinding])])
extractElems Supply
uniqs0 InScopeSet
is1 DataCon
consCon Type
aTy Char
'D' Integer
n Term
arg
            (Either TyVar Type
_ltv:Right Type
snTy:[Either TyVar Type]
_,Type
_) = Type -> ([Either TyVar Type], Type)
splitFunForallTy (TyConMap -> Term -> Type
termType TyConMap
tcm Term
fun)
            (TyConApp TyConName
snatTcNm [Type]
_) = Type -> TypeView
tyView Type
snTy
            (Just TyCon
snatTc)         = TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
lookupUniqMap TyConName
snatTcNm TyConMap
tcm
            [DataCon
snatDc]              = TyCon -> [DataCon]
tyConDataCons TyCon
snatTc
            lbody :: Term
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 :: Term
lb    = [LetBinding] -> Term -> Term
Letrec ([LetBinding] -> [LetBinding]
forall a. [a] -> [a]
init [LetBinding]
elems) Term
lbody
        (Supply -> Identity Supply)
-> RewriteState extra -> Identity (RewriteState extra)
forall extra. Lens' (RewriteState extra) Supply
uniqSupply ((Supply -> Identity Supply)
 -> RewriteState extra -> Identity (RewriteState extra))
-> Supply -> RewriteMonad extra ()
forall s (m :: Type -> Type) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
Lens..= Supply
uniqs1
        Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed Term
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
$ $([Char]
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
  :: InScopeSet
  -> Integer  -- ^ Length of the vector
  -> Type -- ^ Element type of the vector
  -> Term -- ^ The argument vector
  -> NormalizeSession Term
reduceHead :: InScopeSet -> Integer -> Type -> Term -> NormalizeSession Term
reduceHead InScopeSet
inScope Integer
n Type
aTy Term
vArg = do
    TyConMap
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
Lens' RewriteEnv TyConMap
tcCache
    let ty :: Type
ty = TyConMap -> Term -> Type
termType TyConMap
tcm Term
vArg
    TyConMap -> Type -> NormalizeSession Term
forall extra. TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
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
lookupUniqMap TyConName
vecTcNm TyConMap
tcm
      , TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
vecTcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"Clash.Sized.Vector.Vec"
      , [DataCon
_,DataCon
consCon]  <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      = do
        Supply
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. Lens' (RewriteState extra) Supply
uniqSupply
        let (Supply
uniqs1,([Term]
vars,[LetBinding]
elems)) = ([(Term, [LetBinding])] -> ([Term], [LetBinding]))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (([[LetBinding]] -> [LetBinding])
-> ([Term], [[LetBinding]]) -> ([Term], [LetBinding])
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second [[LetBinding]] -> [LetBinding]
forall (t :: Type -> Type) a. Foldable t => t [a] -> [a]
concat (([Term], [[LetBinding]]) -> ([Term], [LetBinding]))
-> ([(Term, [LetBinding])] -> ([Term], [[LetBinding]]))
-> [(Term, [LetBinding])]
-> ([Term], [LetBinding])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Term, [LetBinding])] -> ([Term], [[LetBinding]])
forall a b. [(a, b)] -> ([a], [b])
unzip)
                                  ((Supply, [(Term, [LetBinding])])
 -> (Supply, ([Term], [LetBinding])))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall a b. (a -> b) -> a -> b
$ Supply
-> InScopeSet
-> DataCon
-> Type
-> Char
-> Integer
-> Term
-> (Supply, [(Term, [LetBinding])])
extractElems Supply
uniqs0 InScopeSet
inScope DataCon
consCon Type
aTy Char
'H' Integer
n Term
vArg
            lb :: Term
lb = [LetBinding] -> Term -> Term
Letrec [[LetBinding] -> LetBinding
forall a. [a] -> a
head [LetBinding]
elems] ([Term] -> Term
forall a. [a] -> a
head [Term]
vars)
        (Supply -> Identity Supply)
-> RewriteState extra -> Identity (RewriteState extra)
forall extra. Lens' (RewriteState extra) Supply
uniqSupply ((Supply -> Identity Supply)
 -> RewriteState extra -> Identity (RewriteState extra))
-> Supply -> RewriteMonad extra ()
forall s (m :: Type -> Type) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
Lens..= Supply
uniqs1
        Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed Term
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
$ $([Char]
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
  :: InScopeSet
  -> Integer  -- ^ Length of the vector
  -> Type -- ^ Element type of the vector
  -> Term -- ^ The argument vector
  -> NormalizeSession Term
reduceTail :: InScopeSet -> Integer -> Type -> Term -> NormalizeSession Term
reduceTail InScopeSet
inScope Integer
n Type
aTy Term
vArg = do
    TyConMap
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
Lens' RewriteEnv TyConMap
tcCache
    let ty :: Type
ty = TyConMap -> Term -> Type
termType TyConMap
tcm Term
vArg
    TyConMap -> Type -> NormalizeSession Term
forall extra. TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
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
lookupUniqMap TyConName
vecTcNm TyConMap
tcm
      , TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
vecTcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"Clash.Sized.Vector.Vec"
      , [DataCon
_,DataCon
consCon]  <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      = do
        Supply
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. Lens' (RewriteState extra) Supply
uniqSupply
        let (Supply
uniqs1,([Term]
_,[LetBinding]
elems)) = ([(Term, [LetBinding])] -> ([Term], [LetBinding]))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (([[LetBinding]] -> [LetBinding])
-> ([Term], [[LetBinding]]) -> ([Term], [LetBinding])
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second [[LetBinding]] -> [LetBinding]
forall (t :: Type -> Type) a. Foldable t => t [a] -> [a]
concat (([Term], [[LetBinding]]) -> ([Term], [LetBinding]))
-> ([(Term, [LetBinding])] -> ([Term], [[LetBinding]]))
-> [(Term, [LetBinding])]
-> ([Term], [LetBinding])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Term, [LetBinding])] -> ([Term], [[LetBinding]])
forall a b. [(a, b)] -> ([a], [b])
unzip)
                               ((Supply, [(Term, [LetBinding])])
 -> (Supply, ([Term], [LetBinding])))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall a b. (a -> b) -> a -> b
$ Supply
-> InScopeSet
-> DataCon
-> Type
-> Char
-> Integer
-> Term
-> (Supply, [(Term, [LetBinding])])
extractElems Supply
uniqs0 InScopeSet
inScope DataCon
consCon Type
aTy Char
'L' Integer
n Term
vArg
            b :: LetBinding
b@(Id
tB,Term
_)     = [LetBinding]
elems [LetBinding] -> Unique -> LetBinding
forall a. [a] -> Unique -> a
!! Unique
1
            lb :: Term
lb           = [LetBinding] -> Term -> Term
Letrec [LetBinding
b] (Id -> Term
Var Id
tB)
        (Supply -> Identity Supply)
-> RewriteState extra -> Identity (RewriteState extra)
forall extra. Lens' (RewriteState extra) Supply
uniqSupply ((Supply -> Identity Supply)
 -> RewriteState extra -> Identity (RewriteState extra))
-> Supply -> RewriteMonad extra ()
forall s (m :: Type -> Type) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
Lens..= Supply
uniqs1
        Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed Term
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
$ $([Char]
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
  :: InScopeSet
  -> Integer  -- ^ Length of the vector
  -> Type -- ^ Element type of the vector
  -> Term -- ^ The argument vector
  -> NormalizeSession Term
reduceLast :: InScopeSet -> Integer -> Type -> Term -> NormalizeSession Term
reduceLast InScopeSet
inScope Integer
n Type
aTy Term
vArg = do
    TyConMap
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
Lens' RewriteEnv TyConMap
tcCache
    let ty :: Type
ty = TyConMap -> Term -> Type
termType TyConMap
tcm Term
vArg
    TyConMap -> Type -> NormalizeSession Term
forall extra. TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
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
lookupUniqMap TyConName
vecTcNm TyConMap
tcm
      , TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
vecTcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"Clash.Sized.Vector.Vec"
      , [DataCon
_,DataCon
consCon]  <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      = do
        Supply
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. Lens' (RewriteState extra) Supply
uniqSupply
        let (Supply
uniqs1,([Term]
_,[[LetBinding]]
elems)) = ([(Term, [LetBinding])] -> ([Term], [[LetBinding]]))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [[LetBinding]]))
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second [(Term, [LetBinding])] -> ([Term], [[LetBinding]])
forall a b. [(a, b)] -> ([a], [b])
unzip
                               ((Supply, [(Term, [LetBinding])])
 -> (Supply, ([Term], [[LetBinding]])))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [[LetBinding]]))
forall a b. (a -> b) -> a -> b
$ Supply
-> InScopeSet
-> DataCon
-> Type
-> Char
-> Integer
-> Term
-> (Supply, [(Term, [LetBinding])])
extractElems Supply
uniqs0 InScopeSet
inScope DataCon
consCon Type
aTy Char
'L' Integer
n Term
vArg
            (Id
tB,Term
_)       = [LetBinding] -> LetBinding
forall a. [a] -> a
head ([[LetBinding]] -> [LetBinding]
forall a. [a] -> a
last [[LetBinding]]
elems)
        (Supply -> Identity Supply)
-> RewriteState extra -> Identity (RewriteState extra)
forall extra. Lens' (RewriteState extra) Supply
uniqSupply ((Supply -> Identity Supply)
 -> RewriteState extra -> Identity (RewriteState extra))
-> Supply -> RewriteMonad extra ()
forall s (m :: Type -> Type) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
Lens..= Supply
uniqs1
        case Integer
n of
         Integer
0 -> Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed (Type -> Term
undefinedTm Type
aTy)
         Integer
_ -> Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed ([LetBinding] -> Term -> Term
Letrec ([LetBinding] -> [LetBinding]
forall a. [a] -> [a]
init ([[LetBinding]] -> [LetBinding]
forall (t :: Type -> Type) a. Foldable t => t [a] -> [a]
concat [[LetBinding]]
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
$ $([Char]
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
  :: InScopeSet
  -> PrimInfo -- ^ Primitive info for 'init'
  -> Integer  -- ^ Length of the vector
  -> Type -- ^ Element type of the vector
  -> Term -- ^ The argument vector
  -> NormalizeSession Term
reduceInit :: InScopeSet
-> PrimInfo -> Integer -> Type -> Term -> NormalizeSession Term
reduceInit InScopeSet
_inScope PrimInfo
initPrimInfo Integer
n Type
aTy Term
vArg = do
  TyConMap
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
Lens' RewriteEnv TyConMap
tcCache
  let ty :: Type
ty = TyConMap -> Term -> Type
termType TyConMap
tcm Term
vArg
  Term -> NormalizeSession Term
forall a extra. a -> RewriteMonad extra a
changed (TyConMap -> Type -> Term
go TyConMap
tcm Type
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
lookupUniqMap TyConName
vecTcNm TyConMap
tcm
    , TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
vecTcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"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
$ $([Char]
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
  :: InScopeSet
  -> 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
  -> NormalizeSession Term
reduceAppend :: InScopeSet
-> Integer
-> Integer
-> Type
-> Term
-> Term
-> NormalizeSession Term
reduceAppend InScopeSet
inScope Integer
n Integer
m Type
aTy Term
lArg Term
rArg = do
    TyConMap
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
Lens' RewriteEnv TyConMap
tcCache
    let ty :: Type
ty = TyConMap -> Term -> Type
termType TyConMap
tcm Term
lArg
    TyConMap -> Type -> NormalizeSession Term
forall extra. TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
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
lookupUniqMap TyConName
vecTcNm TyConMap
tcm
      , TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
vecTcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"Clash.Sized.Vector.Vec"
      , [DataCon
_,DataCon
consCon]  <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      = do Supply
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. Lens' (RewriteState extra) Supply
uniqSupply
           let (Supply
uniqs1,([Term]
vars,[LetBinding]
elems)) = ([(Term, [LetBinding])] -> ([Term], [LetBinding]))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (([[LetBinding]] -> [LetBinding])
-> ([Term], [[LetBinding]]) -> ([Term], [LetBinding])
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second [[LetBinding]] -> [LetBinding]
forall (t :: Type -> Type) a. Foldable t => t [a] -> [a]
concat (([Term], [[LetBinding]]) -> ([Term], [LetBinding]))
-> ([(Term, [LetBinding])] -> ([Term], [[LetBinding]]))
-> [(Term, [LetBinding])]
-> ([Term], [LetBinding])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Term, [LetBinding])] -> ([Term], [[LetBinding]])
forall a b. [(a, b)] -> ([a], [b])
unzip)
                                     ((Supply, [(Term, [LetBinding])])
 -> (Supply, ([Term], [LetBinding])))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall a b. (a -> b) -> a -> b
$ Supply
-> InScopeSet
-> DataCon
-> Type
-> Char
-> Integer
-> Term
-> (Supply, [(Term, [LetBinding])])
extractElems Supply
uniqs0 InScopeSet
inScope DataCon
consCon Type
aTy
                                         Char
'C' Integer
n Term
lArg
               lbody :: Term
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) [Term]
vars
               lb :: Term
lb           = [LetBinding] -> Term -> Term
Letrec ([LetBinding] -> [LetBinding]
forall a. [a] -> [a]
init [LetBinding]
elems) Term
lbody
           (Supply -> Identity Supply)
-> RewriteState extra -> Identity (RewriteState extra)
forall extra. Lens' (RewriteState extra) Supply
uniqSupply ((Supply -> Identity Supply)
 -> RewriteState extra -> Identity (RewriteState extra))
-> Supply -> RewriteMonad extra ()
forall s (m :: Type -> Type) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
Lens..= Supply
uniqs1
           Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed Term
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
$ $([Char]
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 :: InScopeSet
               -> PrimInfo -- ^ Unconcat primitive info
               -> Integer  -- ^ Length of the result vector
               -> Integer  -- ^ Length of the elements of the result vector
               -> Type -- ^ Element type
               -> Term -- ^ SNat "Length of the elements of the result vector"
               -> Term -- ^ Argument vector
               -> NormalizeSession Term
reduceUnconcat :: InScopeSet
-> PrimInfo
-> Integer
-> Integer
-> Type
-> Term
-> Term
-> NormalizeSession Term
reduceUnconcat InScopeSet
inScope PrimInfo
unconcatPrimInfo Integer
n Integer
m Type
aTy Term
sm Term
arg = do
    TyConMap
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
Lens' RewriteEnv TyConMap
tcCache
    let ty :: Type
ty = TyConMap -> Term -> Type
termType TyConMap
tcm Term
arg
    TyConMap -> Type -> NormalizeSession Term
forall extra. TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
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
lookupUniqMap TyConName
vecTcNm TyConMap
tcm
      , TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
vecTcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"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 (Unique -> Term -> [Term]
forall a. Unique -> a -> [a]
replicate (Integer -> Unique
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
          Supply
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. Lens' (RewriteState extra) Supply
uniqSupply
          let
            (Supply
uniqs1,([Term]
vars,[LetBinding]
headsAndTails)) =
              ([(Term, [LetBinding])] -> ([Term], [LetBinding]))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (([[LetBinding]] -> [LetBinding])
-> ([Term], [[LetBinding]]) -> ([Term], [LetBinding])
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second [[LetBinding]] -> [LetBinding]
forall (t :: Type -> Type) a. Foldable t => t [a] -> [a]
concat (([Term], [[LetBinding]]) -> ([Term], [LetBinding]))
-> ([(Term, [LetBinding])] -> ([Term], [[LetBinding]]))
-> [(Term, [LetBinding])]
-> ([Term], [LetBinding])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Term, [LetBinding])] -> ([Term], [[LetBinding]])
forall a b. [(a, b)] -> ([a], [b])
unzip)
                     (Supply
-> InScopeSet
-> DataCon
-> Type
-> Char
-> Integer
-> Term
-> (Supply, [(Term, [LetBinding])])
extractElems Supply
uniqs0 InScopeSet
inScope DataCon
consCon Type
aTy Char
'U' (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
*Integer
m) Term
arg)
            -- Build a vector out of the first m elements
            mvec :: Term
mvec = DataCon -> DataCon -> Type -> Integer -> [Term] -> Term
mkVec DataCon
nilCon DataCon
consCon Type
aTy Integer
m (Unique -> [Term] -> [Term]
forall a. Unique -> [a] -> [a]
take (Integer -> Unique
forall a. Num a => Integer -> a
fromInteger Integer
m) [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
            ([LetBinding]
lbs,[LetBinding] -> LetBinding
forall a. [a] -> a
head -> LetBinding
nextVec) = Unique -> [LetBinding] -> ([LetBinding], [LetBinding])
forall a. Unique -> [a] -> ([a], [a])
splitAt ((Unique
2Unique -> Unique -> Unique
forall a. Num a => a -> a -> a
*Integer -> Unique
forall a. Num a => Integer -> a
fromInteger Integer
m)Unique -> Unique -> Unique
forall a. Num a => a -> a -> a
-Unique
1) [LetBinding]
headsAndTails
            -- recursively call unconcat
            nextUnconcat :: Term
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 (LetBinding -> Term
forall a b. (a, b) -> b
snd LetBinding
nextVec)
                                  ]
            -- let (mvec,nextVec) = splitAt sm arg
            -- in Cons mvec (unconcat sm nextVec)
            lBody :: Term
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 :: Term
lb = [LetBinding] -> Term -> Term
Letrec [LetBinding]
lbs Term
lBody

          (Supply -> Identity Supply)
-> RewriteState extra -> Identity (RewriteState extra)
forall extra. Lens' (RewriteState extra) Supply
uniqSupply ((Supply -> Identity Supply)
 -> RewriteState extra -> Identity (RewriteState extra))
-> Supply -> RewriteMonad extra ()
forall s (m :: Type -> Type) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
Lens..= Supply
uniqs1
          Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed Term
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
$ $([Char]
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 -- ^ Argument vector
                -> NormalizeSession Term
reduceTranspose :: Integer -> Integer -> Type -> Term -> NormalizeSession Term
reduceTranspose Integer
n Integer
0 Type
aTy Term
arg = do
    TyConMap
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
Lens' RewriteEnv TyConMap
tcCache
    let ty :: Type
ty = TyConMap -> Term -> Type
termType TyConMap
tcm Term
arg
    TyConMap -> Type -> NormalizeSession Term
forall extra. TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
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
lookupUniqMap TyConName
vecTcNm TyConMap
tcm
      , TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
vecTcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"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 (Unique -> Term -> [Term]
forall a. Unique -> a -> [a]
replicate (Integer -> Unique
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
$ $([Char]
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
_ = [Char] -> NormalizeSession Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> NormalizeSession Term)
-> [Char] -> NormalizeSession Term
forall a b. (a -> b) -> a -> b
$ $([Char]
curLoc) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"reduceTranspose: unimplemented"

reduceReplicate :: Integer
                -> Type
                -> Type
                -> Term
                -> NormalizeSession Term
reduceReplicate :: Integer -> Type -> Type -> Term -> NormalizeSession Term
reduceReplicate Integer
n Type
aTy Type
eTy Term
arg = do
    TyConMap
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
Lens' RewriteEnv TyConMap
tcCache
    TyConMap -> Type -> NormalizeSession Term
forall extra. TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
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
lookupUniqMap TyConName
vecTcNm TyConMap
tcm
      , TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
vecTcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"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 (Unique -> Term -> [Term]
forall a. Unique -> a -> [a]
replicate (Integer -> Unique
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
$ $([Char]
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
  :: InScopeSet
  -> Integer
  -- ^ Size of vector
  -> Type
  -- ^ Type of vector element
  -> Type
  -- ^ Type of vector
  -> Term
  -- ^ Vector
  -> Term
  -- ^ Index
  -> Term
  -- ^ Element
  -> NormalizeSession Term
reduceReplace_int :: InScopeSet
-> Integer
-> Type
-> Type
-> Term
-> Term
-> Term
-> NormalizeSession Term
reduceReplace_int InScopeSet
is0 Integer
n Type
aTy Type
vTy Term
v Term
i Term
newA = do
  TyConMap
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
Lens' RewriteEnv TyConMap
tcCache
  TyConMap -> Type -> NormalizeSession Term
forall extra. TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
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
    (Just TyCon
boolTc) = Unique -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
lookupUniqMap (Unique -> Unique
getKey Unique
boolTyConKey) TyConMap
tcm
    [DataCon
_,DataCon
trueDc]    = TyCon -> [DataCon]
tyConDataCons TyCon
boolTc
    eqInt :: Term
eqInt         = Type -> Type -> Term
eqIntPrim Type
iTy (TyConName -> [Type] -> Type
mkTyConApp (TyCon -> TyConName
tyConName TyCon
boolTc) [])
    case0 :: Term
case0         = Term -> Type -> [Alt] -> Term
Case (Term -> [Either Term Type] -> Term
mkApps Term
eqInt [Term -> Either Term Type
forall a b. a -> Either a b
Left Term
i
                                       ,Term -> Either Term Type
forall a b. a -> Either a b
Left (Term -> [Either Term Type] -> Term
mkApps (DataCon -> Term
Data DataCon
iDc)
                                             [Term -> Either Term Type
forall a b. a -> Either a b
Left (Literal -> Term
Literal (Integer -> Literal
IntLiteral Integer
elIndex))])
                                       ])
                         Type
aTy
                         [(Pat
DefaultPat, Term
oldA)
                         ,(DataCon -> [TyVar] -> [Id] -> Pat
DataPat DataCon
trueDc [] [], Term
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 (OccName -> Type -> WorkInfo -> IsMultiPrim -> PrimInfo
PrimInfo
           OccName
"Clash.Transformations.eqInt"
           (Type -> Type -> Type
mkFunTy Type
intTy (Type -> Type -> Type
mkFunTy Type
intTy Type
boolTy))
           WorkInfo
WorkVariable
           IsMultiPrim
SingleResult )

  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
lookupUniqMap TyConName
vecTcNm TyConMap
tcm
      , TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
vecTcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"Clash.Sized.Vector.Vec"
    , [DataCon
nilCon,DataCon
consCon] <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
    = do
      -- Get data constructors of 'Int'
      Supply
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. Lens' (RewriteState extra) Supply
uniqSupply
      let iTy :: Type
iTy                   = TyConMap -> Term -> Type
termType TyConMap
tcm Term
i
          (TyConApp TyConName
iTcNm [Type]
_)    = Type -> TypeView
tyView Type
iTy
          (Just TyCon
iTc)            = TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
lookupUniqMap TyConName
iTcNm TyConMap
tcm
          [DataCon
iDc]                 = TyCon -> [DataCon]
tyConDataCons TyCon
iTc

      -- Get elements from vector
          (Supply
uniqs1,([Term]
vars,[LetBinding]
elems)) = ([(Term, [LetBinding])] -> ([Term], [LetBinding]))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (([[LetBinding]] -> [LetBinding])
-> ([Term], [[LetBinding]]) -> ([Term], [LetBinding])
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second [[LetBinding]] -> [LetBinding]
forall (t :: Type -> Type) a. Foldable t => t [a] -> [a]
concat (([Term], [[LetBinding]]) -> ([Term], [LetBinding]))
-> ([(Term, [LetBinding])] -> ([Term], [[LetBinding]]))
-> [(Term, [LetBinding])]
-> ([Term], [LetBinding])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Term, [LetBinding])] -> ([Term], [[LetBinding]])
forall a b. [(a, b)] -> ([a], [b])
unzip)
                                ((Supply, [(Term, [LetBinding])])
 -> (Supply, ([Term], [LetBinding])))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall a b. (a -> b) -> a -> b
$ Supply
-> InScopeSet
-> DataCon
-> Type
-> Char
-> Integer
-> Term
-> (Supply, [(Term, [LetBinding])])
extractElems
                                    Supply
uniqs0
                                    InScopeSet
is0
                                    DataCon
consCon
                                    Type
aTy
                                    Char
'I'
                                    Integer
n
                                    Term
v

      -- Replace every element with (if i == elIndex then newA else oldA)
      let replacedEls :: [Term]
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) [Term]
vars [Integer
0..]
          lbody :: Term
lbody       = DataCon -> DataCon -> Type -> Integer -> [Term] -> Term
mkVec DataCon
nilCon DataCon
consCon Type
aTy Integer
n [Term]
replacedEls
          lb :: Term
lb          = [LetBinding] -> Term -> Term
Letrec ([LetBinding] -> [LetBinding]
forall a. [a] -> [a]
init [LetBinding]
elems) Term
lbody
      (Supply -> Identity Supply)
-> RewriteState extra -> Identity (RewriteState extra)
forall extra. Lens' (RewriteState extra) Supply
uniqSupply ((Supply -> Identity Supply)
 -> RewriteState extra -> Identity (RewriteState extra))
-> Supply -> RewriteMonad extra ()
forall s (m :: Type -> Type) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
Lens..= Supply
uniqs1
      Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed Term
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
$ $([Char]
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
  :: InScopeSet
  -> Integer
  -- ^ Size of vector
  -> Type
  -- ^ Type of vector element
  -> Term
  -- ^ Vector
  -> Term
  -- ^ Index
  -> NormalizeSession Term
reduceIndex_int :: InScopeSet
-> Integer -> Type -> Term -> Term -> NormalizeSession Term
reduceIndex_int InScopeSet
is0 Integer
n Type
aTy Term
v Term
i = do
  TyConMap
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
Lens' RewriteEnv TyConMap
tcCache
  let vTy :: Type
vTy = TyConMap -> Term -> Type
termType TyConMap
tcm Term
v
  TyConMap -> Type -> NormalizeSession Term
forall extra. TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
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
    (Just TyCon
boolTc) = Unique -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
lookupUniqMap (Unique -> Unique
getKey Unique
boolTyConKey) TyConMap
tcm
    [DataCon
_,DataCon
trueDc]    = TyCon -> [DataCon]
tyConDataCons TyCon
boolTc
    eqInt :: Term
eqInt         = Type -> Type -> Term
eqIntPrim Type
iTy (TyConName -> [Type] -> Type
mkTyConApp (TyCon -> TyConName
tyConName TyCon
boolTc) [])
    case0 :: Term
case0         = Term -> Type -> [Alt] -> Term
Case (Term -> [Either Term Type] -> Term
mkApps Term
eqInt [Term -> Either Term Type
forall a b. a -> Either a b
Left Term
i
                                       ,Term -> Either Term Type
forall a b. a -> Either a b
Left (Term -> [Either Term Type] -> Term
mkApps (DataCon -> Term
Data DataCon
iDc)
                                             [Term -> Either Term Type
forall a b. a -> Either a b
Left (Literal -> Term
Literal (Integer -> Literal
IntLiteral Integer
elIndex))])
                                       ])
                         Type
aTy
                         [(Pat
DefaultPat, Term
next)
                         ,(DataCon -> [TyVar] -> [Id] -> Pat
DataPat DataCon
trueDc [] [], Term
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 ( OccName -> Type -> WorkInfo -> IsMultiPrim -> PrimInfo
PrimInfo
            OccName
"Clash.Transformations.eqInt"
            (Type -> Type -> Type
mkFunTy Type
intTy (Type -> Type -> Type
mkFunTy Type
intTy Type
boolTy))
            WorkInfo
WorkVariable
            IsMultiPrim
SingleResult )

  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
lookupUniqMap TyConName
vecTcNm TyConMap
tcm
    , TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
vecTcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"Clash.Sized.Vector.Vec"
    , [DataCon
_nilCon,DataCon
consCon] <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
    = do
      -- Get data constructors of 'Int'
      Supply
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. Lens' (RewriteState extra) Supply
uniqSupply
      let iTy :: Type
iTy                   = TyConMap -> Term -> Type
termType TyConMap
tcm Term
i
          (TyConApp TyConName
iTcNm [Type]
_)    = Type -> TypeView
tyView Type
iTy
          (Just TyCon
iTc)            = TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
lookupUniqMap TyConName
iTcNm TyConMap
tcm
          [DataCon
iDc]                 = TyCon -> [DataCon]
tyConDataCons TyCon
iTc

      -- Get elements from vector
          (Supply
uniqs1,([Term]
vars,[LetBinding]
elems)) = ([(Term, [LetBinding])] -> ([Term], [LetBinding]))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (([[LetBinding]] -> [LetBinding])
-> ([Term], [[LetBinding]]) -> ([Term], [LetBinding])
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second [[LetBinding]] -> [LetBinding]
forall (t :: Type -> Type) a. Foldable t => t [a] -> [a]
concat (([Term], [[LetBinding]]) -> ([Term], [LetBinding]))
-> ([(Term, [LetBinding])] -> ([Term], [[LetBinding]]))
-> [(Term, [LetBinding])]
-> ([Term], [LetBinding])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Term, [LetBinding])] -> ([Term], [[LetBinding]])
forall a b. [(a, b)] -> ([a], [b])
unzip)
                                ((Supply, [(Term, [LetBinding])])
 -> (Supply, ([Term], [LetBinding])))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall a b. (a -> b) -> a -> b
$ Supply
-> InScopeSet
-> DataCon
-> Type
-> Char
-> Integer
-> Term
-> (Supply, [(Term, [LetBinding])])
extractElems
                                    Supply
uniqs0
                                    InScopeSet
is0
                                    DataCon
consCon
                                    Type
aTy
                                    Char
'I'
                                    Integer
n
                                    Term
v

      -- Build a right-biased tree of case-expressions
      let indexed :: Term
indexed = ((Term, Integer) -> Term -> Term)
-> Term -> [(Term, Integer)] -> Term
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)
                              (Type -> Term
undefinedTm Type
aTy)
                              ([Term] -> [Integer] -> [(Term, Integer)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Term]
vars [Integer
0..])
          lb :: Term
lb      = [LetBinding] -> Term -> Term
Letrec ([LetBinding] -> [LetBinding]
forall a. [a] -> [a]
init [LetBinding]
elems) Term
indexed
      (Supply -> Identity Supply)
-> RewriteState extra -> Identity (RewriteState extra)
forall extra. Lens' (RewriteState extra) Supply
uniqSupply ((Supply -> Identity Supply)
 -> RewriteState extra -> Identity (RewriteState extra))
-> Supply -> RewriteMonad extra ()
forall s (m :: Type -> Type) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
Lens..= Supply
uniqs1
      Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed Term
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
$ $([Char]
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
  :: InScopeSet
  -> Integer  -- ^ Length of the vector
  -> Type     -- ^ Element type of the argument vector
  -> Term     -- ^ Function to convert elements with
  -> Term     -- ^ Function to combine branches with
  -> Term     -- ^ The vector to fold
  -> NormalizeSession Term
reduceDTFold :: InScopeSet
-> Integer -> Type -> Term -> Term -> Term -> NormalizeSession Term
reduceDTFold InScopeSet
inScope Integer
n Type
aTy Term
lrFun Term
brFun Term
arg = do
    TyConMap
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
Lens' RewriteEnv TyConMap
tcCache
    let ty :: Type
ty = TyConMap -> Term -> Type
termType TyConMap
tcm Term
arg
    TyConMap -> Type -> NormalizeSession Term
forall extra. TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
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
lookupUniqMap TyConName
vecTcNm TyConMap
tcm
      , TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
vecTcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"Clash.Sized.Vector.Vec"
      , [DataCon
_,DataCon
consCon]  <- TyCon -> [DataCon]
tyConDataCons TyCon
vecTc
      = do Supply
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. Lens' (RewriteState extra) Supply
uniqSupply
           let (Supply
uniqs1,([Term]
vars,[LetBinding]
elems)) = ([(Term, [LetBinding])] -> ([Term], [LetBinding]))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (([[LetBinding]] -> [LetBinding])
-> ([Term], [[LetBinding]]) -> ([Term], [LetBinding])
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second [[LetBinding]] -> [LetBinding]
forall (t :: Type -> Type) a. Foldable t => t [a] -> [a]
concat (([Term], [[LetBinding]]) -> ([Term], [LetBinding]))
-> ([(Term, [LetBinding])] -> ([Term], [[LetBinding]]))
-> [(Term, [LetBinding])]
-> ([Term], [LetBinding])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Term, [LetBinding])] -> ([Term], [[LetBinding]])
forall a b. [(a, b)] -> ([a], [b])
unzip)
                                     ((Supply, [(Term, [LetBinding])])
 -> (Supply, ([Term], [LetBinding])))
-> (Supply, [(Term, [LetBinding])])
-> (Supply, ([Term], [LetBinding]))
forall a b. (a -> b) -> a -> b
$ Supply
-> InScopeSet
-> DataCon
-> Type
-> Char
-> Integer
-> Term
-> (Supply, [(Term, [LetBinding])])
extractElems Supply
uniqs0 InScopeSet
inScope DataCon
consCon Type
aTy
                                         Char
'T' (Integer
2Integer -> Integer -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^Integer
n) Term
arg
               (Either TyVar Type
_ltv:Right Type
snTy:[Either TyVar Type]
_,Type
_) = Type -> ([Either TyVar Type], Type)
splitFunForallTy (TyConMap -> Term -> Type
termType TyConMap
tcm Term
brFun)
               (TyConApp TyConName
snatTcNm [Type]
_) = Type -> TypeView
tyView Type
snTy
               (Just TyCon
snatTc)         = TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
lookupUniqMap TyConName
snatTcNm TyConMap
tcm
               [DataCon
snatDc]              = TyCon -> [DataCon]
tyConDataCons TyCon
snatTc
               lbody :: Term
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 :: Term
lb    = [LetBinding] -> Term -> Term
Letrec ([LetBinding] -> [LetBinding]
forall a. [a] -> [a]
init [LetBinding]
elems) Term
lbody
           (Supply -> Identity Supply)
-> RewriteState extra -> Identity (RewriteState extra)
forall extra. Lens' (RewriteState extra) Supply
uniqSupply ((Supply -> Identity Supply)
 -> RewriteState extra -> Identity (RewriteState extra))
-> Supply -> RewriteMonad extra ()
forall s (m :: Type -> Type) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
Lens..= Supply
uniqs1
           Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed Term
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
$ $([Char]
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) = Unique -> [Term] -> ([Term], [Term])
forall a. Unique -> [a] -> ([a], [a])
splitAt (Unique
2Unique -> Integer -> Unique
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
  :: InScopeSet
  -> Integer -- ^ Depth of the tree
  -> Type    -- ^ Element type of the argument tree
  -> Term    -- ^ Function to convert elements with
  -> Term    -- ^ Function to combine branches with
  -> Term    -- ^ The tree to fold
  -> NormalizeSession Term
reduceTFold :: InScopeSet
-> Integer -> Type -> Term -> Term -> Term -> NormalizeSession Term
reduceTFold InScopeSet
inScope Integer
n Type
aTy Term
lrFun Term
brFun Term
arg = do
    TyConMap
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
Lens' RewriteEnv TyConMap
tcCache
    let ty :: Type
ty = TyConMap -> Term -> Type
termType TyConMap
tcm Term
arg
    TyConMap -> Type -> NormalizeSession Term
forall extra. TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
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
lookupUniqMap TyConName
treeTcNm TyConMap
tcm
      , TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
treeTcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"Clash.Sized.RTree.RTree"
      , [DataCon
lrCon,DataCon
brCon] <- TyCon -> [DataCon]
tyConDataCons TyCon
treeTc
      = do Supply
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. Lens' (RewriteState extra) Supply
uniqSupply
           let (Supply
uniqs1,([Term]
vars,[LetBinding]
elems)) = Supply
-> InScopeSet
-> DataCon
-> DataCon
-> Type
-> Char
-> Integer
-> Term
-> (Supply, ([Term], [LetBinding]))
extractTElems Supply
uniqs0 InScopeSet
inScope DataCon
lrCon DataCon
brCon Type
aTy Char
'T' Integer
n Term
arg
               (Either TyVar Type
_ltv:Right Type
snTy:[Either TyVar Type]
_,Type
_) = Type -> ([Either TyVar Type], Type)
splitFunForallTy (TyConMap -> Term -> Type
termType TyConMap
tcm Term
brFun)
               (TyConApp TyConName
snatTcNm [Type]
_) = Type -> TypeView
tyView Type
snTy
               (Just TyCon
snatTc)         = TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
lookupUniqMap TyConName
snatTcNm TyConMap
tcm
               [DataCon
snatDc]              = TyCon -> [DataCon]
tyConDataCons TyCon
snatTc
               lbody :: Term
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 :: Term
lb    = [LetBinding] -> Term -> Term
Letrec [LetBinding]
elems Term
lbody
           (Supply -> Identity Supply)
-> RewriteState extra -> Identity (RewriteState extra)
forall extra. Lens' (RewriteState extra) Supply
uniqSupply ((Supply -> Identity Supply)
 -> RewriteState extra -> Identity (RewriteState extra))
-> Supply -> RewriteMonad extra ()
forall s (m :: Type -> Type) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
Lens..= Supply
uniqs1
           Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed Term
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
$ $([Char]
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) = Unique -> [Term] -> ([Term], [Term])
forall a. Unique -> [a] -> ([a], [a])
splitAt ([Term] -> Unique
forall (t :: Type -> Type) a. Foldable t => t a -> Unique
length [Term]
xs Unique -> Unique -> Unique
forall a. Integral a => a -> a -> a
`div` Unique
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    -- ^ Element
                 -> NormalizeSession Term
reduceTReplicate :: Integer -> Type -> Type -> Term -> NormalizeSession Term
reduceTReplicate Integer
n Type
aTy Type
eTy Term
arg = do
    TyConMap
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
Lens' RewriteEnv TyConMap
tcCache
    TyConMap -> Type -> NormalizeSession Term
forall extra. TyConMap -> Type -> RewriteMonad extra Term
go TyConMap
tcm Type
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
lookupUniqMap TyConName
treeTcNm TyConMap
tcm
      , TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
treeTcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"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 (Unique -> Term -> [Term]
forall a. Unique -> a -> [a]
replicate (Unique
2Unique -> Integer -> Unique
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
$ $([Char]
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)))
         ]