Better fold D{1,2} primitives

1 files changed, 14 insertions, 13 deletions

diff --git a/src/Analysis/Identity.hs b/src/Analysis/Identity.hs
index 2fa156d..9dc8811 100644
--- a/src/Analysis/Identity.hs
+++ b/src/Analysis/Identity.hs
@@ -255,20 +255,21 @@ idana env expr = case expr of
     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
+  EFold1InnerD2 _ cm ef ez eplus ebog ed -> do
+    let STArr _ tB = typeOf ebog
+        t2 = typeOf ez
+    xf1 <- genIds t2
+    xf2 <- genIds tB
+    (_, e1') <- idana (xf1 `SCons` xf2 `SCons` env) ef
+    (_, e2') <- idana env ez
+    xp1 <- genIds t2
+    xp2 <- genIds t2
+    (_, e3') <- idana (xp1 `SCons` xp2 `SCons` env) eplus
+    (v4, 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')
+    let VIArr _ sh = v4
+    res <- VIPair <$> genIds t2 <*> (VIArr <$> genId <*> pure sh)
+    pure (res, EFold1InnerD2 res cm e1' e2' e3' e4' e5')
 
   EConst _ t val -> do
     res <- VIScal <$> genId