From 4d456e4d34b1e4fb3725051d1b8a0c376b704692 Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Thu, 30 Oct 2025 15:56:35 +0100
Subject: Implement reshape

---
 src/AST/Count.hs | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

(limited to 'src/AST/Count.hs')

diff --git a/src/AST/Count.hs b/src/AST/Count.hs
index 229661f..66b4e0b 100644
--- a/src/AST/Count.hs
+++ b/src/AST/Count.hs
@@ -598,6 +598,21 @@ occCountX initialS topexpr k = case topexpr of
   EMaximum1Inner _ e -> handleReduction (EMaximum1Inner ext) e
   EMinimum1Inner _ e -> handleReduction (EMinimum1Inner ext) e
 
+  EReshape _ n esh e ->
+    case s of
+      SsNone ->
+        occCountX SsNone esh $ \env1 mkesh ->
+        occCountX SsNone e $ \env2 mke ->
+        withSome (Some env1 <> Some env2) $ \env ->
+        k env $ \env' ->
+          use (mkesh env') $ use (mke env') $ ENil ext
+      SsArr' s' ->
+        occCountX SsFull esh $ \env1 mkesh ->
+        occCountX (SsArr s') e $ \env2 mke ->
+        withSome (Some env1 <> Some env2) $ \env ->
+        k env $ \env' ->
+          EReshape ext n (mkesh env') (mke env')
+
   EFold1InnerD1 _ cm e1 e2 e3 ->
     case s of
       -- If nothing is necessary, we can execute a fold and then proceed to ignore it
-- 
cgit v1.2.3-70-g09d2