mirror of
https://github.com/anoma/juvix.git
synced 2024-12-13 19:49:20 +03:00
5efec1a9d3
The `rmap` recursor allows to specify changes in binders while going downward through a Core node. This should help in implementing transformations on Core which need to add/remove/change binders. * Depends on PR #1875 * Adds unit tests for `rmap` * Changes the `NatToInt` transformation to use `rmap`
97 lines
3.7 KiB
Haskell
97 lines
3.7 KiB
Haskell
module Core.Recursor.RMap where
|
|
|
|
import Base
|
|
import Core.Recursor.Base
|
|
import Juvix.Compiler.Core.Extra
|
|
import Juvix.Compiler.Core.Language
|
|
|
|
allTests :: TestTree
|
|
allTests =
|
|
testGroup
|
|
"JuvixCore recursors rmap tests"
|
|
(map (mkTest . testDescr) tests)
|
|
|
|
tests :: [UnitTest]
|
|
tests =
|
|
[ UnitTest
|
|
"Identity"
|
|
(rmap (\recur -> recur []))
|
|
( map
|
|
(\x -> (x, x))
|
|
[ mkLambda' mkDynamic' (mkLambda' mkDynamic' (mkVar' 1)),
|
|
mkLambda' mkDynamic' (mkLet' mkDynamic' (mkVar' 0) (mkLambda' mkDynamic' (mkApps' (mkLambdas' [mkDynamic', mkDynamic'] (mkVar' 1)) [mkVar' 1, mkVar' 2]))),
|
|
mkLambdas' [mkTypeInteger', mkTypeInteger'] (mkBuiltinApp' OpIntAdd [mkVar' 1, mkVar' 0])
|
|
]
|
|
),
|
|
UnitTest
|
|
"Add lambdas"
|
|
addLambdas
|
|
[ ( mkLambdas' [mkTypeInteger', mkTypeInteger'] (mkBuiltinApp' OpIntAdd [mkVar' 1, mkVar' 0]),
|
|
mkLambdas' [mkDynamic', mkTypeInteger', mkDynamic', mkTypeInteger'] (mkBuiltinApp' OpIntAdd [mkVar' 2, mkVar' 0])
|
|
)
|
|
],
|
|
UnitTest
|
|
"Replace lambdas with lets"
|
|
replaceLambdasWithLets
|
|
[ ( mkLambdas' [mkTypeInteger', mkTypeInteger'] (mkBuiltinApp' OpIntAdd [mkVar' 1, mkVar' 0]),
|
|
mkLets' [(mkTypeInteger', mkConstant' (ConstInteger 0)), (mkTypeInteger', mkConstant' (ConstInteger 1))] (mkBuiltinApp' OpIntAdd [mkVar' 1, mkVar' 0])
|
|
),
|
|
( mkLambda' mkTypeInteger' $ mkLet' mkTypeInteger' (mkVar' 0) $ mkLambda' mkTypeInteger' $ mkBuiltinApp' OpIntAdd [mkVar' 2, mkVar' 0],
|
|
mkLets' [(mkTypeInteger', mkConstant' (ConstInteger 0)), (mkTypeInteger', mkVar' 0), (mkTypeInteger', mkConstant' (ConstInteger 2))] (mkBuiltinApp' OpIntAdd [mkVar' 2, mkVar' 0])
|
|
)
|
|
],
|
|
UnitTest
|
|
"Remove lambdas"
|
|
removeLambdas
|
|
[ ( mkLambdas' [mkTypeInteger', mkTypeInteger', mkTypeInteger', mkTypeInteger'] (mkBuiltinApp' OpIntAdd [mkBuiltinApp' OpIntAdd [mkVar' 3, mkVar' 1], mkVar' 2]),
|
|
mkLambdas' [mkTypeInteger', mkTypeInteger'] (mkBuiltinApp' OpIntAdd [mkBuiltinApp' OpIntAdd [mkVar' 1, mkVar' 0], mkVar' 1])
|
|
)
|
|
]
|
|
]
|
|
where
|
|
addLambdas :: Node -> Node
|
|
addLambdas = rmap go
|
|
where
|
|
go :: ([BinderChange] -> Node -> Node) -> Node -> Node
|
|
go recur node = case node of
|
|
NLam {} -> mkLambda' mkDynamic' (recur [BCAdd 1] node)
|
|
_ -> recur [] node
|
|
|
|
replaceLambdasWithLets :: Node -> Node
|
|
replaceLambdasWithLets = rmapN go
|
|
where
|
|
go :: ([BinderChange] -> Node -> Node) -> Level -> Node -> Node
|
|
go recur k node = case node of
|
|
NLam Lambda {..} ->
|
|
mkLet
|
|
_lambdaInfo
|
|
(over binderType (cont []) _lambdaBinder)
|
|
(mkConstant' (ConstInteger (fromIntegral k)))
|
|
(cont [BCAdd 1, BCRemove (BinderRemove _lambdaBinder (mkVar' 0))] _lambdaBody)
|
|
_ ->
|
|
recur [] node
|
|
where
|
|
cont :: [BinderChange] -> Node -> Node
|
|
cont bcs = go (recur . (bcs ++)) (k + bindersNumFromBinderChange bcs)
|
|
|
|
removeLambdas :: Node -> Node
|
|
removeLambdas = rmap go
|
|
where
|
|
go :: ([BinderChange] -> Node -> Node) -> Node -> Node
|
|
go recur node = case node of
|
|
NLam lam1@(Lambda _ _ (NLam lam2)) ->
|
|
mkLambda
|
|
(lam1 ^. lambdaInfo)
|
|
(over binderType (cont []) (lam1 ^. lambdaBinder))
|
|
( cont
|
|
[ BCKeep (lam1 ^. lambdaBinder),
|
|
BCRemove (BinderRemove (lam2 ^. lambdaBinder) (mkVar' 0))
|
|
]
|
|
(lam2 ^. lambdaBody)
|
|
)
|
|
_ ->
|
|
recur [] node
|
|
where
|
|
cont :: [BinderChange] -> Node -> Node
|
|
cont bcs = go (recur . (bcs ++))
|