Compile: Add some forgotten refcount decrements

1 files changed, 17 insertions, 1 deletions

diff --git a/src/Compile.hs b/src/Compile.hs
index 7e6fd5c..cb8b424 100644
--- a/src/Compile.hs
+++ b/src/Compile.hs
@@ -48,6 +48,8 @@ let array = arrayGenerate (ShNil `ShCons` 10) (\(IxNil `IxCons` i) -> fromIntegr
 
 -- In shape and index arrays, the innermost dimension is on the right (last index).
 
+-- TODO: test that I'm properly incrementing and decrementing refcounts in all required places
+
 
 debug :: Bool
 debug = toEnum 0
@@ -658,6 +660,8 @@ compile' env = \case
                 pure $ SVerbatim $ accvar ++ " += " ++ argname ++ ".buf->xs[" ++ lenname ++ " * " ++ ivar ++ " + " ++ jvar ++ "];"
              ,SAsg (resname ++ ".buf->xs[" ++ ivar ++ "]") (CELit accvar)]
 
+    incrementVarAlways Decrement (typeOf e) argname
+
     return (CELit resname)
 
   EUnit _ e -> do
@@ -691,6 +695,8 @@ compile' env = \case
                pure $ SAsg (resname ++ ".buf->xs[" ++ ivar ++ " * " ++ lenname ++ " + " ++ jvar ++ "]")
                            (CELit (argname ++ ".buf->xs[" ++ ivar ++ "]"))
 
+    incrementVarAlways Decrement (typeOf earg) argname
+
     return (CELit resname)
 
   EMaximum1Inner _ e -> compileExtremum "max" "maximum1i" ">" env e
@@ -743,7 +749,13 @@ compile' env = \case
   ECustom _ _ _ _ earg _ _ e1 e2 -> do
     name1 <- compileAssign "" env e1
     name2 <- compileAssign "" env e2
-    compile' (Const name2 `SCons` Const name1 `SCons` SNil) earg
+    case (incrementVar Decrement (typeOf e1), incrementVar Decrement (typeOf e2)) of
+      (Nothing, Nothing) -> compile' (Const name2 `SCons` Const name1 `SCons` SNil) earg
+      (mfun1, mfun2) -> do
+        name <- compileAssign "" (Const name2 `SCons` Const name1 `SCons` SNil) earg
+        maybe (return ()) ($ name1) mfun1
+        maybe (return ()) ($ name2) mfun2
+        return (CELit name)
 
   EWith _ t e1 e2 -> do
     name1 <- compileAssign "" env e1
@@ -805,6 +817,8 @@ compile' env = \case
     let dest = accumRef t prj (nameacc++".ac") nameidx
     add (typeOf eval) dest nameval
 
+    incrementVarAlways Decrement (typeOf eval) nameval
+
     return $ CEStruct (repSTy STNil) []
 
   EError _ t s -> do
@@ -1044,6 +1058,8 @@ compileExtremum nameBase opName operator env e = do
               ]
            ,SAsg (resname ++ ".buf->xs[" ++ ivar ++ "]") (CELit redvar)]
 
+  incrementVarAlways Decrement (typeOf e) argname
+
   return (CELit resname)
 
 -- | If this returns Nothing, there was nothing to copy because making a simple