Count: Improve counting of tape in EFold1InnerD1

1 files changed, 34 insertions, 14 deletions

diff --git a/src/AST/Count.hs b/src/AST/Count.hs
index cb363a3..d5afb5e 100644
--- a/src/AST/Count.hs
+++ b/src/AST/Count.hs
@@ -65,6 +65,7 @@ instance Occurrence Occ where
 
 
 data Substruc t t' where
+  -- If you add constructors here, do not forget to update the COMPLETE pragmas of any pattern synonyms below
   SsFull :: Substruc t t
   SsNone :: Substruc t TNil
   SsPair :: Substruc a a' -> Substruc b b' -> Substruc (TPair a b) (TPair a' b')
@@ -79,6 +80,11 @@ pattern SsPair' s1 s2 <- ((\case { SsFull -> SsPair SsFull SsFull ; s -> s }) ->
   where SsPair' = SsPair
 {-# COMPLETE SsNone, SsPair', SsEither, SsLEither, SsMaybe, SsArr, SsAccum #-}
 
+pattern SsArr' :: forall n a t'. forall a'. t' ~ TArr n a' => Substruc a a' -> Substruc (TArr n a) t'
+pattern SsArr' s <- ((\case { SsFull -> SsArr SsFull ; s -> s }) -> SsArr s)
+  where SsArr' = SsArr
+{-# COMPLETE SsNone, SsPair, SsEither, SsLEither, SsMaybe, SsArr', SsAccum #-}
+
 instance Semigroup (Some (Substruc t)) where
   Some SsFull <> _ = Some SsFull
   _ <> Some SsFull = Some SsFull
@@ -592,20 +598,34 @@ occCountX initialS topexpr k = case topexpr of
   EMaximum1Inner _ e -> handleReduction (EMaximum1Inner ext) e
   EMinimum1Inner _ e -> handleReduction (EMinimum1Inner ext) e
 
-  EFold1InnerD1 _ cm a b c ->
-    -- TODO: currently maximally pessimistic on usage here; in particular,
-    -- usage tracking of the 'tape' stores can be improved
-    occCountX SsFull a $ \env1_2' mka ->
-    withSome (scaleMany (Some env1_2')) $ \env1_2 ->
-    occEnvPop' env1_2 $ \env1_1 _ ->
-    occEnvPop' env1_1 $ \env1 _ ->
-    occCountX SsFull b $ \env2 mkb ->
-    occCountX SsFull c $ \env3 mkc ->
-    withSome (Some env1 <> Some env2 <> Some env3) $ \env ->
-    k env $ \env' ->
-      projectSmallerSubstruc SsFull s $
-        EFold1InnerD1 ext cm (mka (OccPush (OccPush env' () SsFull) () SsFull))
-                             (mkb env') (mkc env')
+  EFold1InnerD1{} ->
+    -- TODO: currently somewhat pessimistic on usage here
+    case s of
+      SsPair' _ (SsArr' (SsPair' _ sTape)) ->
+        go topexpr sTape  (projectSmallerSubstruc (SsPair SsFull (SsArr (SsPair SsFull sTape ))) s)
+      -- the primed patsyns also consume SsFull, so if the above doesn't match,
+      -- something along the way is SsNone
+      _ ->
+        go topexpr SsNone (projectSmallerSubstruc (SsPair SsFull (SsArr (SsPair SsFull SsNone))) s)
+    where
+      go :: Expr x env (TPair (TArr n t1) (TArr (S n) (TPair t1 tape)))
+         -> Substruc tape tape'
+         -> (forall env'. Ex env' (TPair (TArr n t1) (TArr (S n) (TPair t1 tape'))) -> Ex env' t')
+         -> r
+      go (EFold1InnerD1 _ cm a b c) sTape project =
+        occCountX (SsPair SsFull sTape) a $ \env1_2' mka ->
+        withSome (scaleMany (Some env1_2')) $ \env1_2 ->
+        occEnvPop' env1_2 $ \env1_1 _ ->
+        occEnvPop' env1_1 $ \env1 _ ->
+        occCountX SsFull b $ \env2 mkb ->
+        occCountX SsFull c $ \env3 mkc ->
+        withSome (Some env1 <> Some env2 <> Some env3) $ \env ->
+        k env $ \env' ->
+          -- projectSmallerSubstruc SsFull s $
+          project $
+            EFold1InnerD1 ext cm (mka (OccPush (OccPush env' () SsFull) () SsFull))
+                                 (mkb env') (mkc env')
+      go _ _ _ = error "impossible"
 
   EFold1InnerD2 _ cm t2 ef ep ezi ebog ed ->
     -- TODO: currently very pessimistic on usage here, can at the very least improve tape usage