{-|
  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

  Partially handles:

    * Clash.Sized.Vector.unconcat
    * Clash.Sized.Vector.transpose
-}

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

module Clash.Normalize.PrimitiveReductions where

import qualified Control.Lens                     as Lens
import           Data.List                        (mapAccumR)
import           Data.List.Extra                  (zipEqual)
import qualified Data.Maybe                       as Maybe

import           PrelNames                        (boolTyConKey)
import           Unique                           (getKey)

import           Clash.Core.DataCon               (DataCon)
import           Clash.Core.Literal               (Literal (..))
import           Clash.Core.Name                  (nameOcc)
import           Clash.Core.Pretty                (showPpr)
import           Clash.Core.Term
  (CoreContext (..), PrimInfo (..), Term (..), WorkInfo (..), Pat (..),
   collectTermIds, mkApps, idToVar)
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)
import           Clash.Core.Util
  (appendToVec, extractElems, extractTElems, mkRTree,
   mkUniqInternalId, mkUniqSystemTyVar, mkVec, dataConInstArgTys,
   primCo, undefinedTm)
import           Clash.Core.Var                   (Var (..))
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

-- | 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
  -> Integer  -- ^ Length of the vector(s)
  -> Type -- ^ Type of the lhs of the function
  -> Type -- ^ Type of the rhs of the function
  -> Type -- ^ 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
-> Integer
-> Type
-> Type
-> Type
-> Term
-> Term
-> Term
-> NormalizeSession Term
reduceZipWith (TransformContext InScopeSet
is0 Context
ctx) 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
    let ty :: Type
ty = TyConMap -> Term -> Type
termType TyConMap
tcm Term
lhsArg
    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, Int, Int) -> CoreContext
AppArg Maybe (OccName, Int, Int)
forall a. Maybe a
NothingCoreContext -> Context -> Context
forall a. a -> [a] -> [a]
:Context
ctx)) Term
fun
        let is1 :: InScopeSet
is1 = InScopeSet -> [Var Term] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 (Term -> [Var Term]
collectTermIds Term
fun1)
            (Supply
uniqs1,([Term]
varsL,[LetBinding]
elemsL)) = ([(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
lhsElTy Char
'L' Integer
n Term
lhsArg
            is2 :: InScopeSet
is2 = InScopeSet -> [Var Term] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 ((LetBinding -> Var Term) -> [LetBinding] -> [Var Term]
forall a b. (a -> b) -> [a] -> [b]
map LetBinding -> Var Term
forall a b. (a, b) -> a
fst [LetBinding]
elemsL)
            (Supply
uniqs2,([Term]
varsR,[LetBinding]
elemsR)) = ([(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
uniqs1 InScopeSet
is2 DataCon
consCon Type
rhsElTy Char
'R' Integer
n Term
rhsArg
            funApps :: [Term]
funApps          = (Term -> Term -> Term) -> [Term] -> [Term] -> [Term]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Term
l Term
r -> Term -> [Either Term Type] -> Term
mkApps Term
fun1 [Term -> Either Term Type
forall a b. a -> Either a b
Left Term
l,Term -> Either Term Type
forall a b. a -> Either a b
Left Term
r]) [Term]
varsL [Term]
varsR
            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]
elemsL [LetBinding] -> [LetBinding] -> [LetBinding]
forall a. [a] -> [a] -> [a]
++ [LetBinding] -> [LetBinding]
forall a. [a] -> [a]
init [LetBinding]
elemsR) 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]
"reduceZipWith: 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.map@ primitive on vectors
-- of a known length @n@, by the fully unrolled recursive "definition" of
-- @Clash.Sized.Vector.map@
reduceMap
  :: TransformContext
  -> 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
-> Integer -> Type -> Type -> Term -> Term -> NormalizeSession Term
reduceMap (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, Int, Int) -> CoreContext
AppArg Maybe (OccName, Int, Int)
forall a. Maybe a
NothingCoreContext -> Context -> Context
forall a. a -> [a] -> [a]
:Context
ctx)) Term
fun
        let is1 :: InScopeSet
is1 = InScopeSet -> [Var Term] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 (Term -> [Var Term]
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
argElTy Char
'A' 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            = 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
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]
"reduceMap: 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.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, Int, Int) -> CoreContext
AppArg Maybe (OccName, Int, Int)
forall a. Maybe a
NothingCoreContext -> Context -> Context
forall a. a -> [a] -> [a]
:Context
ctx)) Term
fun
        let is1 :: InScopeSet
is1 = InScopeSet -> [Var Term] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 (Term -> [Var Term]
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 -> PrimInfo
PrimInfo OccName
"Clash.Sized.Internal.Index.fromInteger#" Type
idxFromIntegerTy WorkInfo
WorkNever)
            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.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, Int, Int) -> CoreContext
AppArg Maybe (OccName, Int, Int)
forall a. Maybe a
NothingCoreContext -> Context -> Context
forall a. a -> [a] -> [a]
:Context
ctx)) Term
fun
        let is1 :: InScopeSet
is1 = InScopeSet -> [Var Term] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 (Term -> [Var Term]
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 :: [Var Term]
apDictIds@[Var Term
functorDictId,Var Term
pureId,Var Term
apId,Var Term
_,Var Term
_,Var Term
_]) =
              ((Supply, InScopeSet)
 -> (OccName, Type) -> ((Supply, InScopeSet), Var Term))
-> (Supply, InScopeSet)
-> [(OccName, Type)]
-> ((Supply, InScopeSet), [Var Term])
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), Var Term)
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 :: [Var Term]
funcDicIds@[Var Term
fmapId,Var Term
_]) =
              ((Supply, InScopeSet)
 -> (OccName, Type) -> ((Supply, InScopeSet), Var Term))
-> (Supply, InScopeSet)
-> [(OccName, Type)]
-> ((Supply, InScopeSet), [Var Term])
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), Var Term)
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] -> [Var Term] -> Pat
DataPat DataCon
apDictCon   [] [Var Term]
apDictIds
            fnPat :: Pat
fnPat    = DataCon -> [TyVar] -> [Var Term] -> Pat
DataPat DataCon
funcDictCon [] [Var Term]
funcDicIds

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

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

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

            -- Extract the 'fmap' function from the Functor dictionary
            fmapTy :: Type
fmapTy = Var Term -> Type
forall a. Var a -> Type
varType Var Term
fmapId
            fmapTm :: Term
fmapTm = Term -> Type -> [Alt] -> Term
Case (Var Term -> Term
Var Var Term
functorDictId) Type
fmapTy
                          [(Pat
fnPat, Var Term -> Term
Var Var Term
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 (Var Term -> Term
Var ([Var Term]
apDictIds[Var Term] -> Int -> Var Term
forall a. [a] -> Int -> a
!!Int
1))
                                                       (Var Term -> Term
Var ([Var Term]
apDictIds[Var Term] -> Int -> Var Term
forall a. [a] -> Int -> a
!!Int
2))
                                                       (Var Term -> Term
Var ([Var Term]
funcDicIds[Var Term] -> Int -> Var Term
forall a. [a] -> Int -> a
!!Int
0))
                                                       Type
bTy Integer
n [Term]
funApps

            lb :: Term
lb      = [LetBinding] -> Term -> Term
Letrec ([(([Var Term]
apDictIds[Var Term] -> Int -> Var Term
forall a. [a] -> Int -> a
!!Int
0), Term
funcTm)
                              ,(([Var Term]
apDictIds[Var Term] -> Int -> Var Term
forall a. [a] -> Int -> a
!!Int
1), Term
pureTm)
                              ,(([Var Term]
apDictIds[Var Term] -> Int -> Var Term
forall a. [a] -> Int -> a
!!Int
2), Term
apTm)
                              ,(([Var Term]
funcDicIds[Var Term] -> Int -> Var Term
forall a. [a] -> Int -> a
!!Int
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
  -> 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
-> Integer -> Type -> Term -> Term -> Term -> NormalizeSession Term
reduceFoldr TransformContext
_ Integer
0 Type
_ Term
_ Term
start Term
_ = Term -> NormalizeSession Term
forall a extra. a -> RewriteMonad extra a
changed Term
start
reduceFoldr (TransformContext InScopeSet
is0 Context
ctx) 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
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, Int, Int) -> CoreContext
AppArg Maybe (OccName, Int, Int)
forall a. Maybe a
NothingCoreContext -> Context -> Context
forall a. a -> [a] -> [a]
:Context
ctx)) Term
fun
        let is1 :: InScopeSet
is1 = InScopeSet -> [Var Term] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 (Term -> [Var Term]
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
'G' Integer
n Term
arg
            lbody :: Term
lbody            = (Term -> Term -> Term) -> Term -> [Term] -> Term
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Term
l Term
r -> Term -> [Either Term Type] -> Term
mkApps Term
fun1 [Term -> Either Term Type
forall a b. a -> Either a b
Left Term
l,Term -> Either Term Type
forall a b. a -> Either a b
Left Term
r]) Term
start [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]
"reduceFoldr: 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.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, Int, Int) -> CoreContext
AppArg Maybe (OccName, Int, Int)
forall a. Maybe a
NothingCoreContext -> Context -> Context
forall a. a -> [a] -> [a]
:Context
ctx)) Term
fun
        let is1 :: InScopeSet
is1 = InScopeSet -> [Var Term] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 (Term -> [Var Term]
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) = Int -> [Term] -> ([Term], [Term])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Term] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [Term]
as Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) [Term]
as
                      lF :: Term
lF    = Term -> [Term] -> Term
foldV Term
f [Term]
l
                      rF :: Term
rF    = Term -> [Term] -> Term
foldV Term
f [Term]
r
                  in  Term -> [Either Term Type] -> Term
mkApps Term
f [Term -> Either Term Type
forall a b. a -> Either a b
Left Term
lF, Term -> Either Term Type
forall a b. a -> Either a b
Left Term
rF]

-- | Replace an application of the @Clash.Sized.Vector.dfold@ primitive on
-- vectors of a known length @n@, by the fully unrolled recursive "definition"
-- of @Clash.Sized.Vector.dfold@
reduceDFold
  :: 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 -> [Var Term] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 (Term -> [Var Term]
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@(Var Term
tB,Term
_)     = [LetBinding]
elems [LetBinding] -> Int -> LetBinding
forall a. [a] -> Int -> a
!! Int
1
            lb :: Term
lb           = [LetBinding] -> Term -> Term
Letrec [LetBinding
b] (Var Term -> Term
Var Var Term
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
            (Var Term
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)) (Var Term -> Term
Var Var Term
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
  -> Integer  -- ^ Length of the vector
  -> Type -- ^ Element type of the vector
  -> Term -- ^ The argument vector
  -> NormalizeSession Term
reduceInit :: InScopeSet -> Integer -> Type -> Term -> NormalizeSession Term
reduceInit 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
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]
_,[[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
        (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
1 -> Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed (DataCon -> DataCon -> Type -> Integer -> [Term] -> Term
mkVec DataCon
nilCon DataCon
consCon Type
aTy Integer
0 [])
         Integer
_ -> let el :: [[LetBinding]]
el = [[LetBinding]] -> [[LetBinding]]
forall a. [a] -> [a]
init [[LetBinding]]
elems
                  iv :: Term
iv = DataCon -> DataCon -> Type -> Integer -> [Term] -> Term
mkVec DataCon
nilCon DataCon
consCon Type
aTy (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1) (([LetBinding] -> Term) -> [[LetBinding]] -> [Term]
forall a b. (a -> b) -> [a] -> [b]
map (Var Term -> Term
idToVar (Var Term -> Term)
-> ([LetBinding] -> Var Term) -> [LetBinding] -> Term
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LetBinding -> Var Term
forall a b. (a, b) -> a
fst (LetBinding -> Var Term)
-> ([LetBinding] -> LetBinding) -> [LetBinding] -> Var Term
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [LetBinding] -> LetBinding
forall a. [a] -> a
head) [[LetBinding]]
el)
                  lb :: [LetBinding]
lb = [LetBinding] -> [LetBinding]
forall a. [a] -> [a]
init ([[LetBinding]] -> [LetBinding]
forall (t :: Type -> Type) a. Foldable t => t [a] -> [a]
concat [[LetBinding]]
el)
              in  Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed ([LetBinding] -> Term -> Term
Letrec [LetBinding]
lb Term
iv)

    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]
"reduceInit: 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.(++)@ 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 :: Integer  -- ^ Length of the result vector
               -> Integer  -- ^ Length of the elements of the result vector
               -> Type -- ^ Element type
               -> Term -- ^ Argument vector
               -> NormalizeSession Term
reduceUnconcat :: Integer -> Integer -> Type -> Term -> NormalizeSession Term
reduceUnconcat 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 (Int -> Term -> [Term]
forall a. Int -> a -> [a]
replicate (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
n) Term
nilVec)
        in  Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed Term
retVec
    go TyConMap
_ Type
ty = [Char] -> RewriteMonad extra Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> RewriteMonad extra Term)
-> [Char] -> RewriteMonad extra Term
forall a b. (a -> b) -> a -> b
$ $([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

reduceUnconcat 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]
"reduceUnconcat: unimplemented"

-- | 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 (Int -> Term -> [Term]
forall a. Int -> a -> [a]
replicate (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
n) Term
nilVec)
        in  Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed Term
retVec
    go TyConMap
_ Type
ty = [Char] -> RewriteMonad extra Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> RewriteMonad extra Term)
-> [Char] -> RewriteMonad extra Term
forall a b. (a -> b) -> a -> b
$ $([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 (Int -> Term -> [Term]
forall a. Int -> a -> [a]
replicate (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
n) Term
arg)
        in  Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed Term
retVec
    go TyConMap
_ Type
ty = [Char] -> RewriteMonad extra Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> RewriteMonad extra Term)
-> [Char] -> RewriteMonad extra Term
forall a b. (a -> b) -> a -> b
$ $([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) = Int -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
lookupUniqMap (Unique -> Int
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] -> [Var Term] -> 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 -> PrimInfo
PrimInfo OccName
"Clash.Transformations.eqInt" (Type -> Type -> Type
mkFunTy Type
intTy (Type -> Type -> Type
mkFunTy Type
intTy Type
boolTy)) WorkInfo
WorkVariable)

  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) = Int -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
lookupUniqMap (Unique -> Int
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] -> [Var Term] -> 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 -> PrimInfo
PrimInfo OccName
"Clash.Transformations.eqInt" (Type -> Type -> Type
mkFunTy Type
intTy (Type -> Type -> Type
mkFunTy Type
intTy Type
boolTy)) WorkInfo
WorkVariable)

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

-- | Replace an application of the @Clash.Sized.RTree.tdfold@ primitive on
-- trees of a known depth @n@, by the fully unrolled recursive "definition"
-- of @Clash.Sized.RTree.tdfold@
reduceTFold
  :: 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) = Int -> [Term] -> ([Term], [Term])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Term] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [Term]
xs Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) [Term]
xs
          k' :: Integer
k'        = Integer
kInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1
          eL :: Term
eL        = (Integer -> Term) -> Integer -> [Term] -> Term
doFold Integer -> Term
snDc Integer
k' [Term]
xsL
          eR :: Term
eR        = (Integer -> Term) -> Integer -> [Term] -> Term
doFold Integer -> Term
snDc Integer
k' [Term]
xsR
      in  Term -> [Either Term Type] -> Term
mkApps Term
brFun [Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
k))
                       ,Term -> Either Term Type
forall a b. a -> Either a b
Left (Integer -> Term
snDc Integer
k)
                       ,Term -> Either Term Type
forall a b. a -> Either a b
Left Term
eL
                       ,Term -> Either Term Type
forall a b. a -> Either a b
Left Term
eR
                       ]

reduceTReplicate :: Integer -- ^ Depth of the tree
                 -> Type    -- ^ Element type
                 -> Type    -- ^ Result type
                 -> Term    -- ^ 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 (Int -> Term -> [Term]
forall a. Int -> a -> [a]
replicate (Int
2Int -> Integer -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^Integer
n) Term
arg)
        in  Term -> RewriteMonad extra Term
forall a extra. a -> RewriteMonad extra a
changed Term
retVec
    go TyConMap
_ Type
ty = [Char] -> RewriteMonad extra Term
forall a. HasCallStack => [Char] -> a
error ([Char] -> RewriteMonad extra Term)
-> [Char] -> RewriteMonad extra Term
forall a b. (a -> b) -> a -> b
$ $([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)))
         ]