Compile: Decrement array elements when freeing array

1 files changed, 15 insertions, 4 deletions

diff --git a/src/Compile.hs b/src/Compile.hs
index cb8b424..b9dbd41 100644
--- a/src/Compile.hs
+++ b/src/Compile.hs
@@ -860,7 +860,7 @@ incrementVar inc ty =
 incrementVarAlways :: Increment -> STy a -> String -> CompM ()
 incrementVarAlways inc ty var = maybe (pure ()) ($ var) (incrementVar inc ty)
 
-data ArrayTree = ATArray  -- ^ we've arrived at an array we need to decrement the refcount of
+data ArrayTree = ATArray (Some SNat) (Some STy)  -- ^ we've arrived at an array we need to decrement the refcount of (contains rank and element type of the array)
                | ATNoop  -- ^ don't do anything here
                | ATProj String ArrayTree  -- ^ descend one field deeper
                | ATCondTag ArrayTree ArrayTree  -- ^ if tag is 0, first; if 1, second
@@ -886,16 +886,27 @@ makeArrayTree (STPair a b) = smartATBoth (smartATProj "a" (makeArrayTree a))
 makeArrayTree (STEither a b) = smartATCondTag (smartATProj "l" (makeArrayTree a))
                                               (smartATProj "r" (makeArrayTree b))
 makeArrayTree (STMaybe t) = smartATCondTag ATNoop (makeArrayTree t)
-makeArrayTree (STArr _ _) = ATArray
+makeArrayTree (STArr n t) = ATArray (Some n) (Some t)
 makeArrayTree (STScal _) = ATNoop
 makeArrayTree (STAccum _) = ATNoop
 
 incrementVar' :: Increment -> String -> ArrayTree -> CompM ()
-incrementVar' inc path ATArray =
+incrementVar' inc path (ATArray (Some n) (Some eltty)) =
   case inc of
     Increment -> emit $ SVerbatim (path ++ ".buf->refc++;")
     Decrement ->
-      emit $ SVerbatim $ "if (--" ++ path ++ ".buf->refc == 0) free(" ++ path ++ ".buf);"
+      case incrementVar Decrement eltty of
+        Nothing -> emit $ SVerbatim $ "if (--" ++ path ++ ".buf->refc == 0) free(" ++ path ++ ".buf);"
+        Just f -> do
+          shszvar <- genName' "frshsz"
+          ivar <- genName' "i"
+          ((), eltDecrStmts) <- scope $ f (path ++ ".buf->xs[" ++ ivar ++ "]")
+          emit $ SIf (CELit ("--" ++ path ++ ".buf->refc == 0"))
+                   (BList [SVarDecl True "size_t" shszvar (compileArrShapeSize n path)
+                          ,SLoop "size_t" ivar (CELit "0") (CELit shszvar) $
+                             BList eltDecrStmts
+                          ,SVerbatim $ "free(" ++ path ++ ".buf);"])
+                   mempty
 incrementVar' _ _ ATNoop = pure ()
 incrementVar' inc path (ATProj field t) = incrementVar' inc (path ++ "." ++ field) t
 incrementVar' inc path (ATCondTag t1 t2) = do