From b8c5c85e24df82462f03b47b298573206e7b7f80 Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Wed, 8 Oct 2025 20:51:37 +0200
Subject: Count: Correct alternative and scaleMany counting

---
 src/AST/Count.hs | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

(limited to 'src')

diff --git a/src/AST/Count.hs b/src/AST/Count.hs
index f055164..5666289 100644
--- a/src/AST/Count.hs
+++ b/src/AST/Count.hs
@@ -434,8 +434,8 @@ occCountX initialS topexpr k = case topexpr of
     occEnvPop' env1' $ \env1 s1 ->
     occEnvPop' env2' $ \env2 s2 ->
     occCountX (SsEither s1 s2) e $ \env0 mke ->
-    withSome (Some env0 <> Some env1) $ \env01 ->
-    withSome (Some env01 <> Some env2) $ \env ->
+    withSome (Some env1 <||> Some env2) $ \env12 ->
+    withSome (Some env12 <> Some env0) $ \env ->
     k env $ \env' ->
       ECase ext (mke env') (mka (OccPush env' () s1)) (mkb (OccPush env' () s2))
   ENothing _ t ->
@@ -459,8 +459,8 @@ occCountX initialS topexpr k = case topexpr of
     occCountX s b $ \env2' mkb ->
     occEnvPop' env2' $ \env2 s2 ->
     occCountX (SsMaybe s2) e $ \env0 mke ->
-    withSome (Some env0 <> Some env1) $ \env01 ->
-    withSome (Some env01 <> Some env2) $ \env ->
+    withSome (Some env1 <||> Some env2) $ \env12 ->
+    withSome (Some env12 <> Some env0) $ \env ->
     k env $ \env' ->
       EMaybe ext (mka env') (mkb (OccPush env' () s2)) (mke env')
   ELNil _ t1 t2 ->
@@ -497,9 +497,9 @@ occCountX initialS topexpr k = case topexpr of
     occEnvPop' env2' $ \env2 s1 ->
     occEnvPop' env3' $ \env3 s2 ->
     occCountX (SsLEither s1 s2) e $ \env0 mke ->
-    withSome (Some env0 <> Some env1) $ \env01 ->
-    withSome (Some env01 <> Some env2) $ \env012 ->
-    withSome (Some env012 <> Some env3) $ \env ->
+    withSome (Some env1 <||> Some env2) $ \env12 ->
+    withSome (Some env12 <||> Some env3) $ \env123 ->
+    withSome (Some env123 <> Some env0) $ \env ->
     k env $ \env' ->
       ELCase ext (mke env') (mka env') (mkb (OccPush env' () s1)) (mkc (OccPush env' () s2))
 
@@ -514,7 +514,8 @@ occCountX initialS topexpr k = case topexpr of
     case s of
       SsNone ->
         occCountX SsFull a $ \env1 mka ->
-        occCountX SsNone b $ \env2' mkb ->
+        occCountX SsNone b $ \env2'' mkb ->
+        withSome (scaleMany (Some env2'')) $ \env2' ->
         occEnvPop' env2' $ \env2 s2 ->
         withSome (Some env1 <> Some env2) $ \env ->
         k env $ \env' ->
@@ -525,7 +526,8 @@ occCountX initialS topexpr k = case topexpr of
             ENil ext
       SsArr s' ->
         occCountX SsFull a $ \env1 mka ->
-        occCountX s' b $ \env2' mkb ->
+        occCountX s' b $ \env2'' mkb ->
+        withSome (scaleMany (Some env2'')) $ \env2' ->
         occEnvPop' env2' $ \env2 s2 ->
         withSome (Some env1 <> Some env2) $ \env ->
         k env $ \env' ->
@@ -535,7 +537,8 @@ occCountX initialS topexpr k = case topexpr of
       SsFull -> occCountX (SsArr SsFull) topexpr k
 
   EFold1Inner _ commut a b c ->
-    occCountX SsFull a $ \env1'' mka ->
+    occCountX SsFull a $ \env1''' mka ->
+    withSome (scaleMany (Some env1''')) $ \env1'' ->
     occEnvPop' env1'' $ \env1' s2 ->
     occEnvPop' env1' $ \env1 s1 ->
     let s0 = case s of
-- 
cgit v1.2.3-70-g09d2