WIP fold: everything but Compile (slow, but should be sound)

1 files changed, 26 insertions, 0 deletions

diff --git a/src/Analysis/Identity.hs b/src/Analysis/Identity.hs
index b54946b..2fa156d 100644
--- a/src/Analysis/Identity.hs
+++ b/src/Analysis/Identity.hs
@@ -244,6 +244,32 @@ idana env expr = case expr of
     res <- VIArr <$> genId <*> pure sh
     pure (res, EMinimum1Inner res e1')
 
+  EFold1InnerD1 _ cm e1 e2 e3 -> do
+    let t1 = typeOf e2
+    x1 <- genIds t1
+    x2 <- genIds t1
+    (_, e1') <- idana (x1 `SCons` x2 `SCons` env) e1
+    (_, e2') <- idana env e2
+    (v3, e3') <- idana env e3
+    let VIArr _ sh'@(_ :< sh) = v3
+    res <- VIPair <$> (VIArr <$> genId <*> pure sh) <*> (VIArr <$> genId <*> pure sh')
+    pure (res, EFold1InnerD1 res cm e1' e2' e3')
+
+  EFold1InnerD2 _ cm t2 ef ep ezi ebog ed -> do
+    let STArr _ (STPair t1 ttape) = typeOf ebog
+    x1 <- genIds (fromSMTy t2)
+    x2 <- genIds t1
+    x3 <- genIds t1
+    x4 <- genIds ttape
+    (_, e1') <- idana (x1 `SCons` x2 `SCons` x3 `SCons` x4 `SCons` env) ef
+    (v2, e2') <- idana env ep
+    (_, e3') <- idana env ezi
+    (_, e4') <- idana env ebog
+    (_, e5') <- idana env ed
+    let VIArr _ sh = v2
+    res <- VIPair <$> genIds (fromSMTy t2) <*> (VIArr <$> genId <*> pure sh)
+    pure (res, EFold1InnerD2 res cm t2 e1' e2' e3' e4' e5')
+
   EConst _ t val -> do
     res <- VIScal <$> genId
     pure (res, EConst res t val)