From ddb2b99bef797b11805257eb1a5712a2c4d285a6 Mon Sep 17 00:00:00 2001
From: Tobias Mock <mail@tmock.de>
Date: Fri, 27 Sep 2024 21:34:49 +0200
Subject: [PATCH] Fix double free codegen

when passing parts of parameters in tail calls.
---
 lib/mallocs.ml                   | 18 ++++++++++--
 lib/monomorph_tree.ml            |  8 +-----
 test/misc.t/free_param_parts.smu | 47 ++++++++++++++++++++++++++++++++
 test/misc.t/run.t                |  7 +++++
 4 files changed, 71 insertions(+), 9 deletions(-)
 create mode 100644 test/misc.t/free_param_parts.smu

diff --git a/lib/mallocs.ml b/lib/mallocs.ml
index 06492338..ec0c5d15 100644
--- a/lib/mallocs.ml
+++ b/lib/mallocs.ml
@@ -18,10 +18,24 @@ module Make (Mtree : Monomorph_tree_intf.S) = struct
 
   type pmap = Pset.t Imap.t
 
+  let rec is_arg = function
+    | Malloc.No_malloc -> true
+    | Param _ -> true
+    | Single mid -> is_arg_mid mid
+    | Path (m, _) -> is_arg m
+
+  and is_arg_mid (mid : Mid.t) =
+    match mid.parent with Some m -> is_arg m | None -> false
+
   let mlist_of_pmap m =
     Imap.to_rev_seq m
-    |> Seq.map (fun ((id : Mid.t), paths) ->
-           { id = id.mid; mtyp = id.typ; paths })
+    |> Seq.filter_map (fun ((id : Mid.t), paths) ->
+           (* If the malloc comes from a borrowed parameter, we don't add it to the
+              list of mallocs. This list is later on used for freeing allocs. Since
+              there are paths and parent relationsships, we need to recursively
+              check if something is an argument. *)
+           if is_arg_mid id then None
+           else Some { id = id.mid; mtyp = id.typ; paths })
     |> List.of_seq
 
   let show_pmap m =
diff --git a/lib/monomorph_tree.ml b/lib/monomorph_tree.ml
index 0ab1bbc9..2bfd8ba8 100644
--- a/lib/monomorph_tree.ml
+++ b/lib/monomorph_tree.ml
@@ -1178,12 +1178,6 @@ and morph_app mk p callee args ret_typ =
   in
 
   let f p (arg, attr) =
-    let is_arg = function
-      | Malloc.No_malloc -> false
-      | Param _ -> true
-      | Single _ -> false
-      | Path _ -> (* A path cannot be a passed argument *) false
-    in
     let ret = p.ret in
     let p, ex, var = morph_expr { p with ret = false } arg in
     let is_moved =
@@ -1199,7 +1193,7 @@ and morph_app mk p callee args ret_typ =
       monomorph,
       (* If an argument is passed by move this means the parameter is also owned
          and will be freed in a tailrec call *)
-      is_arg var.malloc || is_moved )
+      Monomorph_impl.Mallocs_ipml.is_arg var.malloc || is_moved )
   in
 
   let rec fold_decr_last p args = function
diff --git a/test/misc.t/free_param_parts.smu b/test/misc.t/free_param_parts.smu
new file mode 100644
index 00000000..e2b6b4c1
--- /dev/null
+++ b/test/misc.t/free_param_parts.smu
@@ -0,0 +1,47 @@
+module typekey {
+  signature {
+  type t
+}
+type t = int
+}
+
+type named = { thing : string, idx : int }
+
+fun rec consume(storage, opt) {
+  match storage.[opt] {
+    #none: println("none")
+    #some(thing, idx): {
+      println(thing)
+      -- passing idx in a tail call caused the string 'thing' to be freed. The
+      -- solution is to explicitly mark it as (child of a) parameter, and not
+      -- adding parameters to free lists.
+      consume(storage, idx)
+    }
+  }
+}
+
+consume([#some("thing", 1), #none], 0)
+
+-- simpler examples without the array/get
+fun rec simpler(arg, _) {
+  match arg {
+    #some(thing, idx): {
+      ignore(thing)
+      simpler(#none, idx)}
+    #none: ()
+  }
+}
+
+simpler(#some("thing", 0), 0)
+
+-- moved version
+fun rec simpler(arg!, _!) {
+  match !arg {
+    #some(thing, idx!): {
+      ignore(thing)
+      simpler(!#none, !idx)}
+    #none: ()
+  }
+}
+
+simpler(!#some("thing", 0), !0)
diff --git a/test/misc.t/run.t b/test/misc.t/run.t
index 66aad68d..220f588a 100644
--- a/test/misc.t/run.t
+++ b/test/misc.t/run.t
@@ -3992,3 +3992,10 @@ Using unit values
 `inner` here should not make `tmp` a const, otherwise could gen would fail
   $ schmu mutable_inner_let.smu
   $ valgrind -q --leak-check=yes --show-reachable=yes ./mutable_inner_let
+
+
+Don't free params if parts are passed in tail calls
+  $ schmu free_param_parts.smu
+  $ valgrind -q --leak-check=yes --show-reachable=yes ./free_param_parts
+  thing
+  none