module CLaSH.Normalize.PrimitiveReductions where
import qualified Control.Lens as Lens
import qualified Data.HashMap.Lazy as HashMap
import qualified Data.Maybe as Maybe
import Data.Text (pack)
import Unbound.Generics.LocallyNameless (bind, embed, rec, rebind,
string2Name, name2String)
import CLaSH.Core.DataCon (DataCon, dataConInstArgTys,
dcName, dcType)
import CLaSH.Core.Literal (Literal (..))
import CLaSH.Core.Term (Term (..), Pat (..))
import CLaSH.Core.Type (LitTy (..), Type (..),
TypeView (..), mkFunTy,
mkTyConApp, splitFunForallTy,
tyView)
import CLaSH.Core.TyCon (TyConName, tyConDataCons)
import CLaSH.Core.TysPrim (typeNatKind)
import CLaSH.Core.Util (appendToVec, extractElems,
idToVar, mkApps, mkVec,
termType)
import CLaSH.Core.Var (Var (..))
import CLaSH.Normalize.Types
import CLaSH.Rewrite.Types
import CLaSH.Rewrite.Util
import CLaSH.Util
reduceZipWith :: Int
-> Type
-> Type
-> Type
-> Term
-> Term
-> Term
-> NormalizeSession Term
reduceZipWith n lhsElTy rhsElTy resElTy fun lhsArg rhsArg = do
tcm <- Lens.view tcCache
(TyConApp vecTcNm _) <- tyView <$> termType tcm lhsArg
let (Just vecTc) = HashMap.lookup vecTcNm tcm
[nilCon,consCon] = tyConDataCons vecTc
(varsL,elemsL) = second concat . unzip
$ extractElems consCon lhsElTy 'L' n lhsArg
(varsR,elemsR) = second concat . unzip
$ extractElems consCon rhsElTy 'R' n rhsArg
funApps = zipWith (\l r -> mkApps fun [Left l,Left r]) varsL varsR
lbody = mkVec nilCon consCon resElTy n funApps
lb = Letrec (bind (rec (init elemsL ++ init elemsR)) lbody)
changed lb
reduceMap :: Int
-> Type
-> Type
-> Term
-> Term
-> NormalizeSession Term
reduceMap n argElTy resElTy fun arg = do
tcm <- Lens.view tcCache
(TyConApp vecTcNm _) <- tyView <$> termType tcm arg
let (Just vecTc) = HashMap.lookup vecTcNm tcm
[nilCon,consCon] = tyConDataCons vecTc
(vars,elems) = second concat . unzip
$ extractElems consCon argElTy 'A' n arg
funApps = map (fun `App`) vars
lbody = mkVec nilCon consCon resElTy n funApps
lb = Letrec (bind (rec (init elems)) lbody)
changed lb
reduceTraverse :: Int
-> Type
-> Type
-> Type
-> Term
-> Term
-> Term
-> NormalizeSession Term
reduceTraverse n aTy fTy bTy dict fun arg = do
tcm <- Lens.view tcCache
(TyConApp vecTcNm _) <- tyView <$> termType tcm arg
(TyConApp apDictTcNm _) <- tyView <$> termType tcm dict
let (Just apDictTc) = HashMap.lookup apDictTcNm tcm
[apDictCon] = tyConDataCons apDictTc
(Just apDictIdTys) = dataConInstArgTys apDictCon [fTy]
apDictIds = zipWith Id (map string2Name ["functorDict"
,"pure"
,"ap"
,"apConstL"
,"apConstR"])
(map embed apDictIdTys)
(TyConApp funcDictTcNm _) = tyView (head apDictIdTys)
(Just funcDictTc) = HashMap.lookup funcDictTcNm tcm
[funcDictCon] = tyConDataCons funcDictTc
(Just funcDictIdTys) = dataConInstArgTys funcDictCon [fTy]
funcDicIds = zipWith Id (map string2Name ["fmap","fmapConst"])
(map embed funcDictIdTys)
apPat = DataPat (embed apDictCon) (rebind [] apDictIds)
fnPat = DataPat (embed funcDictCon) (rebind [] funcDicIds)
pureTy = apDictIdTys!!1
pureTm = Case dict pureTy [bind apPat (Var pureTy (string2Name "pure"))]
apTy = apDictIdTys!!2
apTm = Case dict apTy [bind apPat (Var apTy (string2Name "ap"))]
funcTy = (head apDictIdTys)
funcTm = Case dict funcTy
[bind apPat (Var funcTy (string2Name "functorDict"))]
fmapTy = (head funcDictIdTys)
fmapTm = Case (Var funcTy (string2Name "functorDict")) fmapTy
[bind fnPat (Var fmapTy (string2Name "fmap"))]
(Just vecTc) = HashMap.lookup vecTcNm tcm
[nilCon,consCon] = tyConDataCons vecTc
(vars,elems) = second concat . unzip
$ extractElems consCon aTy 'T' n arg
funApps = map (fun `App`) vars
lbody = mkTravVec vecTcNm nilCon consCon (idToVar (apDictIds!!1))
(idToVar (apDictIds!!2))
(idToVar (funcDicIds!!0))
bTy n funApps
lb = Letrec (bind (rec ([((apDictIds!!0),embed funcTm)
,((apDictIds!!1),embed pureTm)
,((apDictIds!!2),embed apTm)
,((funcDicIds!!0),embed fmapTm)
] ++ init elems)) lbody)
changed lb
mkTravVec :: TyConName
-> DataCon
-> DataCon
-> Term
-> Term
-> Term
-> Type
-> Int
-> [Term]
-> Term
mkTravVec vecTc nilCon consCon pureTm apTm fmapTm bTy = go
where
go :: Int -> [Term] -> Term
go _ [] = mkApps pureTm [Right (mkTyConApp vecTc [LitTy (NumTy 0),bTy])
,Left (mkApps (Data nilCon)
[Right (LitTy (NumTy 0))
,Right bTy
,Left (Prim "_CO_" nilCoTy)])]
go n (x:xs) = mkApps apTm
[Right (mkTyConApp vecTc [LitTy (NumTy (n1)),bTy])
,Right (mkTyConApp vecTc [LitTy (NumTy n),bTy])
,Left (mkApps fmapTm [Right bTy
,Right (mkFunTy (mkTyConApp vecTc [LitTy (NumTy (n1)),bTy])
(mkTyConApp vecTc [LitTy (NumTy n),bTy]))
,Left (mkApps (Data consCon)
[Right (LitTy (NumTy n))
,Right bTy
,Right (LitTy (NumTy (n1)))
,Left (Prim "_CO_" (consCoTy n))
])
,Left x])
,Left (go (n1) xs)]
nilCoTy = head (Maybe.fromJust (dataConInstArgTys nilCon [(LitTy (NumTy 0))
,bTy]))
consCoTy n = head (Maybe.fromJust (dataConInstArgTys consCon
[(LitTy (NumTy n))
,bTy
,(LitTy (NumTy (n1)))]))
reduceFoldr :: Int
-> Type
-> Type
-> Term
-> Term
-> Term
-> NormalizeSession Term
reduceFoldr n aTy _bTy fun start arg = do
tcm <- Lens.view tcCache
(TyConApp vecTcNm _) <- tyView <$> termType tcm arg
let (Just vecTc) = HashMap.lookup vecTcNm tcm
[_,consCon] = tyConDataCons vecTc
(vars,elems) = second concat . unzip
$ extractElems consCon aTy 'G' n arg
lbody = foldr (\l r -> mkApps fun [Left l,Left r]) start vars
lb = Letrec (bind (rec (init elems)) lbody)
changed lb
reduceFold :: Int
-> Type
-> Term
-> Term
-> NormalizeSession Term
reduceFold n aTy fun arg = do
tcm <- Lens.view tcCache
(TyConApp vecTcNm _) <- tyView <$> termType tcm arg
let (Just vecTc) = HashMap.lookup vecTcNm tcm
[_,consCon] = tyConDataCons vecTc
(vars,elems) = second concat . unzip
$ extractElems consCon aTy 'F' n arg
lbody = foldV vars
lb = Letrec (bind (rec (init elems)) lbody)
changed lb
where
foldV [a] = a
foldV as = let (l,r) = splitAt (length as `div` 2) as
lF = foldV l
rF = foldV r
in mkApps fun [Left lF, Left rF]
reduceDFold :: Int
-> Type
-> Term
-> Term
-> Term
-> NormalizeSession Term
reduceDFold n aTy fun start arg = do
tcm <- Lens.view tcCache
(TyConApp vecTcNm _) <- tyView <$> termType tcm arg
let (Just vecTc) = HashMap.lookup vecTcNm tcm
[_,consCon] = tyConDataCons vecTc
(vars,elems) = second concat . unzip
$ extractElems consCon aTy 'D' n arg
([_ltv,Right snTy,_etaTy,_eta1Ty],_) <- splitFunForallTy <$> termType tcm fun
let (TyConApp snatTcNm _) = tyView snTy
(Just snatTc) = HashMap.lookup snatTcNm tcm
[snatDc] = tyConDataCons snatTc
([_nTv,_kn,Right pTy],_) = splitFunForallTy (dcType snatDc)
(TyConApp proxyTcNm _) = tyView pTy
(Just proxyTc) = HashMap.lookup proxyTcNm tcm
[proxyDc] = tyConDataCons proxyTc
buildSNat i = mkApps (Prim (pack (name2String (dcName snatDc)))
(dcType snatDc))
[Right (LitTy (NumTy i))
,Left (Literal (IntegerLiteral (toInteger i)))
,Left (mkApps (Data proxyDc)
[Right typeNatKind
,Right (LitTy (NumTy i))])
]
lbody = doFold buildSNat (n1) vars
lb = Letrec (bind (rec (init elems)) lbody)
changed lb
where
doFold _ _ [] = start
doFold snDc k (x:xs) = mkApps fun
[Right (LitTy (NumTy k))
,Left (snDc k)
,Left x
,Left (doFold snDc (k1) xs)
]
reduceHead :: Int
-> Type
-> Term
-> NormalizeSession Term
reduceHead n aTy vArg = do
tcm <- Lens.view tcCache
(TyConApp vecTcNm _) <- tyView <$> termType tcm vArg
let (Just vecTc) = HashMap.lookup vecTcNm tcm
[_,consCon] = tyConDataCons vecTc
(vars,elems) = second concat . unzip
$ extractElems consCon aTy 'H' n vArg
lb = Letrec (bind (rec [head elems]) (head vars))
changed lb
reduceTail :: Int
-> Type
-> Term
-> NormalizeSession Term
reduceTail n aTy vArg = do
tcm <- Lens.view tcCache
(TyConApp vecTcNm _) <- tyView <$> termType tcm vArg
let (Just vecTc) = HashMap.lookup vecTcNm tcm
[_,consCon] = tyConDataCons vecTc
(_,elems) = second concat . unzip
$ extractElems consCon aTy 'L' n vArg
b@(tB,_) = elems !! 1
lb = Letrec (bind (rec [b]) (idToVar tB))
changed lb
reduceAppend :: Int
-> Int
-> Type
-> Term
-> Term
-> NormalizeSession Term
reduceAppend n m aTy lArg rArg = do
tcm <- Lens.view tcCache
(TyConApp vecTcNm _) <- tyView <$> termType tcm lArg
let (Just vecTc) = HashMap.lookup vecTcNm tcm
[_,consCon] = tyConDataCons vecTc
(vars,elems) = second concat . unzip
$ extractElems consCon aTy 'C' n lArg
lbody = appendToVec consCon aTy rArg (n+m) vars
lb = Letrec (bind (rec (init elems)) lbody)
changed lb
reduceUnconcat :: Int
-> Int
-> Type
-> Term
-> NormalizeSession Term
reduceUnconcat n 0 aTy arg = do
tcm <- Lens.view tcCache
(TyConApp vecTcNm _) <- tyView <$> termType tcm arg
let (Just vecTc) = HashMap.lookup vecTcNm tcm
[nilCon,consCon] = tyConDataCons vecTc
nilVec = mkVec nilCon consCon aTy 0 []
innerVecTy = mkTyConApp vecTcNm [LitTy (NumTy 0), aTy]
retVec = mkVec nilCon consCon innerVecTy n (replicate n nilVec)
changed retVec
reduceUnconcat _ _ _ _ = error $ $(curLoc) ++ "reduceUnconcat: unimplemented"
reduceTranspose :: Int
-> Int
-> Type
-> Term
-> NormalizeSession Term
reduceTranspose n 0 aTy arg = do
tcm <- Lens.view tcCache
(TyConApp vecTcNm _) <- tyView <$> termType tcm arg
let (Just vecTc) = HashMap.lookup vecTcNm tcm
[nilCon,consCon] = tyConDataCons vecTc
nilVec = mkVec nilCon consCon aTy 0 []
innerVecTy = mkTyConApp vecTcNm [LitTy (NumTy 0), aTy]
retVec = mkVec nilCon consCon innerVecTy n (replicate n nilVec)
changed retVec
reduceTranspose _ _ _ _ = error $ $(curLoc) ++ "reduceTranspose: unimplemented"