Add ext field to remaining AST constructors

1 files changed, 35 insertions, 33 deletions

diff --git a/src/Simplify.hs b/src/Simplify.hs
index 6303716..785e2bd 100644
--- a/src/Simplify.hs
+++ b/src/Simplify.hs
@@ -93,33 +93,33 @@ simplify' = \case
   -- TODO: constant folding for operations
 
   -- TODO: properly concatenate accum/onehot
-  EAccum SZ _ (EOneHot _ i idx val) acc ->
+  EAccum _ SZ _ (EOneHot _ _ i idx val) acc ->
     acted $ simplify' $
-      EAccum i idx val acc
-  EAccum _ _ (EZero _) _ -> (Any True, ENil ext)
-  EPlus _ (EZero _) e -> acted $ simplify' e
-  EPlus _ e (EZero _) -> acted $ simplify' e
-  EOneHot _ SZ _ e -> acted $ simplify' e
+      EAccum ext i idx val acc
+  EAccum _ _ _ (EZero _ _) _ -> (Any True, ENil ext)
+  EPlus _ _ (EZero _ _) e -> acted $ simplify' e
+  EPlus _ _ e (EZero _ _) -> acted $ simplify' e
+  EOneHot _ _ SZ _ e -> acted $ simplify' e
 
   -- equations for plus
-  EPlus STNil _ _ -> (Any True, ENil ext)
+  EPlus _ STNil _ _ -> (Any True, ENil ext)
 
-  EPlus (STPair t1 t2) (EJust _ (EPair _ a1 b1)) (EJust _ (EPair _ a2 b2)) ->
-    acted $ simplify' $ EJust ext (EPair ext (EPlus t1 a1 a2) (EPlus t2 b1 b2))
-  EPlus STPair{} ENothing{} e -> acted $ simplify' e
-  EPlus STPair{} e ENothing{} -> acted $ simplify' e
+  EPlus _ (STPair t1 t2) (EJust _ (EPair _ a1 b1)) (EJust _ (EPair _ a2 b2)) ->
+    acted $ simplify' $ EJust ext (EPair ext (EPlus ext t1 a1 a2) (EPlus ext t2 b1 b2))
+  EPlus _ STPair{} ENothing{} e -> acted $ simplify' e
+  EPlus _ STPair{} e ENothing{} -> acted $ simplify' e
 
-  EPlus (STEither t1 _) (EJust _ (EInl _ dt2 a1)) (EJust _ (EInl _ _ a2)) ->
-    acted $ simplify' $ EJust ext (EInl ext dt2 (EPlus t1 a1 a2))
-  EPlus (STEither _ t2) (EJust _ (EInr _ dt1 b1)) (EJust _ (EInr _ _ b2)) ->
-    acted $ simplify' $ EJust ext (EInr ext dt1 (EPlus t2 b1 b2))
-  EPlus STEither{} ENothing{} e -> acted $ simplify' e
-  EPlus STEither{} e ENothing{} -> acted $ simplify' e
+  EPlus _ (STEither t1 _) (EJust _ (EInl _ dt2 a1)) (EJust _ (EInl _ _ a2)) ->
+    acted $ simplify' $ EJust ext (EInl ext dt2 (EPlus ext t1 a1 a2))
+  EPlus _ (STEither _ t2) (EJust _ (EInr _ dt1 b1)) (EJust _ (EInr _ _ b2)) ->
+    acted $ simplify' $ EJust ext (EInr ext dt1 (EPlus ext t2 b1 b2))
+  EPlus _ STEither{} ENothing{} e -> acted $ simplify' e
+  EPlus _ STEither{} e ENothing{} -> acted $ simplify' e
 
-  EPlus (STMaybe t) (EJust _ e1) (EJust _ e2) ->
-    acted $ simplify' $ EJust ext (EPlus t e1 e2)
-  EPlus STMaybe{} ENothing{} e -> acted $ simplify' e
-  EPlus STMaybe{} e ENothing{} -> acted $ simplify' e
+  EPlus _ (STMaybe t) (EJust _ e1) (EJust _ e2) ->
+    acted $ simplify' $ EJust ext (EPlus ext t e1 e2)
+  EPlus _ STMaybe{} ENothing{} e -> acted $ simplify' e
+  EPlus _ STMaybe{} e ENothing{} -> acted $ simplify' e
 
   -- fallback recursion
   EVar _ t i -> pure $ EVar ext t i
@@ -154,12 +154,12 @@ simplify' = \case
       <*> (let ?accumInScope = False in simplify' b)
       <*> (let ?accumInScope = False in simplify' c)
       <*> simplify' e1 <*> simplify' e2
-  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
-  EZero t -> pure $ EZero t
-  EPlus t a b -> EPlus t <$> simplify' a <*> simplify' b
-  EOneHot t i a b -> EOneHot t i <$> simplify' a <*> simplify' b
-  EError t s -> pure $ EError t s
+  EWith _ e1 e2 -> EWith ext <$> simplify' e1 <*> (let ?accumInScope = True in simplify' e2)
+  EAccum _ i e1 e2 e3 -> EAccum ext i <$> simplify' e1 <*> simplify' e2 <*> simplify' e3
+  EZero _ t -> pure $ EZero ext t
+  EPlus _ t a b -> EPlus ext t <$> simplify' a <*> simplify' b
+  EOneHot _ t i a b -> EOneHot ext t i <$> simplify' a <*> simplify' b
+  EError _ t s -> pure $ EError ext t s
 
 acted :: (Any, a) -> (Any, a)
 acted (_, x) = (Any True, x)
@@ -169,6 +169,8 @@ cheapExpr = \case
   EVar{} -> True
   ENil{} -> True
   EConst{} -> True
+  EFst _ e -> cheapExpr e
+  ESnd _ e -> cheapExpr e
   _ -> False
 
 -- | This can be made more precise by tracking (and not counting) adds on
@@ -202,12 +204,12 @@ hasAdds = \case
   EIdx _ a b -> hasAdds a || hasAdds b
   EShape _ e -> hasAdds e
   EOp _ _ e -> hasAdds e
-  EWith a b -> hasAdds a || hasAdds b
-  EAccum _ _ _ _ -> True
-  EZero _ -> False
-  EPlus _ a b -> hasAdds a || hasAdds b
-  EOneHot _ _ a b -> hasAdds a || hasAdds b
-  EError _ _ -> False
+  EWith _ a b -> hasAdds a || hasAdds b
+  EAccum _ _ _ _ _ -> True
+  EZero _ _ -> False
+  EPlus _ _ a b -> hasAdds a || hasAdds b
+  EOneHot _ _ _ a b -> hasAdds a || hasAdds b
+  EError _ _ _ -> False
 
 checkAccumInScope :: SList STy env -> Bool
 checkAccumInScope = \case SNil -> False