blob: f2fc54a9e311f9c5fd81387443cc6e5e3b83b71d (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
|
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ImplicitParams #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
module Simplify where
import Data.Monoid
import AST
import AST.Count
import Data
simplifyN :: KnownEnv env => Int -> Ex env t -> Ex env t
simplifyN 0 = id
simplifyN n = simplifyN (n - 1) . simplify
simplify :: forall env t. KnownEnv env => Ex env t -> Ex env t
simplify = let ?accumInScope = checkAccumInScope @env knownEnv in simplify'
simplify' :: (?accumInScope :: Bool) => Ex env t -> Ex env t
simplify' = \case
-- inlining
ELet _ rhs body
| not ?accumInScope || not (hasAdds rhs) -- cannot discard effectful computations
, Occ lexOcc runOcc <- occCount IZ body
, lexOcc <= One -- prevent code size blowup
, runOcc <= One -- prevent runtime increase
-> simplify' (subst1 rhs body)
| cheapExpr rhs
-> simplify' (subst1 rhs body)
-- let splitting
ELet _ (EPair _ a b) body ->
simplify' $
ELet ext a $
ELet ext (weakenExpr WSink b) $
subst (\_ t -> \case IZ -> EPair ext (EVar ext (typeOf a) (IS IZ)) (EVar ext (typeOf b) IZ)
IS i -> EVar ext t (IS (IS i)))
body
-- let rotation
ELet _ (ELet _ rhs a) b ->
ELet ext (simplify' rhs) $
ELet ext (simplify' a) $
weakenExpr (WCopy WSink) (simplify' b)
-- beta rules for products
EFst _ (EPair _ e _) -> simplify' e
ESnd _ (EPair _ _ e) -> simplify' e
-- beta rules for coproducts
ECase _ (EInl _ _ e) rhs _ -> simplify' (ELet ext e rhs)
ECase _ (EInr _ _ e) _ rhs -> simplify' (ELet ext e rhs)
-- TODO: array indexing (index of build, index of fold)
-- TODO: constant folding for operations
EVar _ t i -> EVar ext t i
ELet _ a b -> ELet ext (simplify' a) (simplify' b)
EPair _ a b -> EPair ext (simplify' a) (simplify' b)
EFst _ e -> EFst ext (simplify' e)
ESnd _ e -> ESnd ext (simplify' e)
ENil _ -> ENil ext
EInl _ t e -> EInl ext t (simplify' e)
EInr _ t e -> EInr ext t (simplify' e)
ECase _ e a b -> ECase ext (simplify' e) (simplify' a) (simplify' b)
EBuild1 _ a b -> EBuild1 ext (simplify' a) (simplify' b)
EBuild _ es e -> EBuild ext (fmap simplify' es) (simplify' e)
EFold1 _ a b -> EFold1 ext (simplify' a) (simplify' b)
EUnit _ e -> EUnit ext (simplify' e)
EConst _ t v -> EConst ext t v
EIdx0 _ e -> EIdx0 ext (simplify' e)
EIdx1 _ a b -> EIdx1 ext (simplify' a) (simplify' b)
EIdx _ e es -> EIdx ext (simplify' e) (fmap simplify' es)
EOp _ op e -> EOp ext op (simplify' e)
EWith e1 e2 -> EWith (simplify' e1) (let ?accumInScope = True in simplify' e2)
EAccum1 e1 e2 e3 -> EAccum1 (simplify' e1) (simplify' e2) (simplify' e3)
EError t s -> EError t s
cheapExpr :: Expr x env t -> Bool
cheapExpr = \case
EVar{} -> True
ENil{} -> True
EConst{} -> True
_ -> False
-- | This can be made more precise by tracking (and not counting) adds on
-- locally eliminated accumulators.
hasAdds :: Expr x env t -> Bool
hasAdds = \case
EVar _ _ _ -> False
ELet _ rhs body -> hasAdds rhs || hasAdds body
EPair _ a b -> hasAdds a || hasAdds b
EFst _ e -> hasAdds e
ESnd _ e -> hasAdds e
ENil _ -> False
EInl _ _ e -> hasAdds e
EInr _ _ e -> hasAdds e
ECase _ e a b -> hasAdds e || hasAdds a || hasAdds b
EBuild1 _ a b -> hasAdds a || hasAdds b
EBuild _ es e -> getAny (foldMap (Any . hasAdds) es) || hasAdds e
EFold1 _ a b -> hasAdds a || hasAdds b
EUnit _ e -> hasAdds e
EConst _ _ _ -> False
EIdx0 _ e -> hasAdds e
EIdx1 _ a b -> hasAdds a || hasAdds b
EIdx _ e es -> hasAdds e || getAny (foldMap (Any . hasAdds) es)
EOp _ _ e -> hasAdds e
EWith a b -> hasAdds a || hasAdds b
EAccum1 _ _ _ -> True
EError _ _ -> False
checkAccumInScope :: SList STy env -> Bool
checkAccumInScope = \case SNil -> False
SCons t env -> check t || checkAccumInScope env
where
check :: STy t -> Bool
check STNil = False
check (STPair s t) = check s || check t
check (STEither s t) = check s || check t
check (STArr _ t) = check t
check (STScal _) = False
check STAccum{} = True
|
