Implement reshape

1 files changed, 15 insertions, 0 deletions

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