blob: c44d965f14333742f4061b41b6758e4ab6ac5e2b (
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
130
131
132
133
134
135
136
137
138
139
140
|
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ImplicitParams #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
module Simplify where
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
| cheapExpr rhs
-> simplify' (subst1 rhs body)
| Occ lexOcc runOcc <- occCount IZ body
, ((not ?accumInScope || not (hasAdds rhs)) && lexOcc <= One && runOcc <= One) -- without effects, normal rules apply
|| (lexOcc == One && runOcc == One) -- with effects, linear inlining is still allowed, but weakening is not
-> 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 ->
simplify' $
ELet ext rhs $
ELet ext 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)
-- let floating to facilitate beta reduction
EFst _ (ELet _ rhs body) -> simplify' (ELet ext rhs (EFst ext body))
ESnd _ (ELet _ rhs body) -> simplify' (ELet ext rhs (ESnd ext body))
ECase _ (ELet _ rhs body) e1 e2 -> simplify' (ELet ext rhs (ECase ext body (weakenExpr (WCopy WSink) e1) (weakenExpr (WCopy WSink) e2)))
EIdx0 _ (ELet _ rhs body) -> simplify' (ELet ext rhs (EIdx0 ext body))
EIdx1 _ (ELet _ rhs body) e -> simplify' (ELet ext rhs (EIdx1 ext body (weakenExpr WSink e)))
-- 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 _ n a b -> EBuild ext n (simplify' a) (simplify' b)
EFold1 _ a b -> EFold1 ext (simplify' a) (simplify' b)
EUnit _ e -> EUnit ext (simplify' e)
-- EReplicate _ e -> EReplicate 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 _ n a b -> EIdx ext n (simplify' a) (simplify' b)
EShape _ e -> EShape ext (simplify' e)
EOp _ op e -> EOp ext op (simplify' e)
EWith e1 e2 -> EWith (simplify' e1) (let ?accumInScope = True in simplify' e2)
EAccum i e1 e2 e3 -> EAccum i (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 _ _ a b -> hasAdds a || hasAdds b
EFold1 _ a b -> hasAdds a || hasAdds b
EUnit _ e -> hasAdds e
-- EReplicate _ e -> hasAdds e
EConst _ _ _ -> False
EIdx0 _ e -> hasAdds e
EIdx1 _ a b -> hasAdds a || hasAdds b
EIdx _ _ a b -> hasAdds a || hasAdds b
EShape _ e -> hasAdds e
EOp _ _ e -> hasAdds e
EWith a b -> hasAdds a || hasAdds b
EAccum _ _ _ _ -> 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
|
