Support partial freeing of rcs

lib/codegen/autogen.ml

@@ -115,7 +115,14 @@ struct
                   |> ignore;
                   decl_children_exc pset pseudovar f.ftyp)
           fields
-    | _ -> failwith "TODO decl free or not supported"
+    | Trc _ ->
+        (* This is a rc where the payload has been moved out and the refcount
+           still needs freeing. There is no need to declare anything, because we
+           won't free the payload. *)
+        ()
+    | t ->
+        print_endline (show_typ t);
+        failwith "TODO decl free or not supported"
 
   let copy param allocref v =
     if contains_allocation v.typ then
@@ -390,6 +397,45 @@ struct
         Hashtbl.add cls_func_tbl name func;
         func
 
+  let free_rc v item_typ =
+    let v = bring_default_var v in
+
+    let rf = Llvm.build_gep int_t v.value [| ci 0 |] "ref" builder in
+
+    let rc = Llvm.build_load int_t rf "refc" builder in
+
+    (* Get current block *)
+    let start_bb = Llvm.insertion_block builder in
+    let parent = Llvm.block_parent start_bb in
+
+    let decr_bb = Llvm.append_block context "decr" parent in
+    let free_bb = Llvm.append_block context "free" parent in
+    let merge_bb = Llvm.append_block context "merge" parent in
+
+    let cmp =
+      Llvm.(build_icmp Icmp.Eq) rc (Llvm.const_int int_t 1) "" builder
+    in
+    ignore (Llvm.build_cond_br cmp free_bb decr_bb builder);
+
+    (* decr *)
+    Llvm.position_at_end decr_bb builder;
+    let added = Llvm.build_sub rc (Llvm.const_int int_t 1) "" builder in
+    ignore (Llvm.build_store added rf builder);
+    ignore (Llvm.build_br merge_bb builder);
+
+    (* free *)
+    Llvm.position_at_end free_bb builder;
+    let value = Llvm.build_gep int_t v.value [| ci 1 |] "vl" builder in
+    (match item_typ with
+    | Some item_typ ->
+        free_call_only
+          { value; typ = item_typ; lltyp = get_lltype_def item_typ; kind = Ptr }
+    | None -> ());
+    free_var v.value |> ignore;
+    ignore (Llvm.build_br merge_bb builder);
+
+    Llvm.position_at_end merge_bb builder
+
   let free_impl v =
     match v.typ with
     | Tarray t ->
@@ -401,41 +447,7 @@ struct
            iter_array_children v sz t free_call_only);
 
         free_var v.value |> ignore
-    | Trc item_typ ->
-        let v = bring_default_var v in
-
-        let rf = Llvm.build_gep int_t v.value [| ci 0 |] "ref" builder in
-
-        let rc = Llvm.build_load int_t rf "refc" builder in
-
-        (* Get current block *)
-        let start_bb = Llvm.insertion_block builder in
-        let parent = Llvm.block_parent start_bb in
-
-        let decr_bb = Llvm.append_block context "decr" parent in
-        let free_bb = Llvm.append_block context "free" parent in
-        let merge_bb = Llvm.append_block context "merge" parent in
-
-        let cmp =
-          Llvm.(build_icmp Icmp.Eq) rc (Llvm.const_int int_t 1) "" builder
-        in
-        ignore (Llvm.build_cond_br cmp free_bb decr_bb builder);
-
-        (* decr *)
-        Llvm.position_at_end decr_bb builder;
-        let added = Llvm.build_sub rc (Llvm.const_int int_t 1) "" builder in
-        ignore (Llvm.build_store added rf builder);
-        ignore (Llvm.build_br merge_bb builder);
-
-        (* free *)
-        Llvm.position_at_end free_bb builder;
-        let value = Llvm.build_gep int_t v.value [| ci 1 |] "vl" builder in
-        free_call_only
-          { value; typ = item_typ; lltyp = get_lltype_def item_typ; kind = Ptr };
-        free_var v.value |> ignore;
-        ignore (Llvm.build_br merge_bb builder);
-
-        Llvm.position_at_end merge_bb builder
+    | Trc item_typ -> free_rc v (Some item_typ)
     | Trecord (_, Rec_folded, _) -> failwith "unreachable"
     | Trecord (_, (Rec_not fields | Rec_top fields), _) ->
         Array.iteri
@@ -537,6 +549,9 @@ struct
                   let v = follow_field v i in
                   free_except_call pset v)
           fields
+    | Trc _ ->
+        (* Only free refcount ptr. Same as free above, but without the payload *)
+        free_rc v None
     | _ -> failwith "TODO free or not supported"
 
   let gen_functions () =

lib/codegen/codegen.ml

@@ -207,7 +207,12 @@ end = struct
           | _ -> gen_app param callee args alloca typed_expr.typ
         in
 
-        List.iter (fun id -> Strtbl.replace free_tbl id value) ms;
+        (match callee.monomorph with
+        | Builtin (Rc_get, _) ->
+            (* Don't add frees. They would overwrite the rc frees. Think record
+               field: We don't add the child frees either. *)
+            ()
+        | _ -> List.iter (fun id -> Strtbl.replace free_tbl id value) ms);
         fin value
     | Mif expr -> gen_if param expr
     | Mrecord (labels, allocref, id) ->
@@ -1053,8 +1058,7 @@ end = struct
           | _ -> failwith "Internal Error: Arity mismatch in builder"
         in
         R.gen_rc param e fnc.ret allocref
-    | Rc_get ->
-        List.hd args |> bring_default_var |> fun llvar -> R.get llvar fnc.ret
+    | Rc_get -> List.hd args |> bring_default_var |> fun llvar -> R.get llvar
     | Any_abort -> (
         let ft, abort =
           Llvm.(
@@ -1450,7 +1454,14 @@ end = struct
     let open Malloc_types in
     let expr = gen_expr param expr in
     let get_path path init =
-      List.fold_right (fun index expr -> follow_field expr index) path init
+      List.fold_right
+        (fun index expr ->
+          match index with
+          | -1 ->
+              (* Special case for rc get *)
+              R.get expr
+          | _ -> follow_field expr index)
+        path init
     in
     (match fs with
     | Except fs ->

lib/codegen/rc.ml

@@ -16,7 +16,7 @@ module type S = sig
     Monomorph_tree.allocas ref ->
     llvar
 
-  val get : llvar -> typ -> llvar
+  val get : llvar -> llvar
 end
 
 module Make (C : Core) (T : Lltypes_intf.S) (H : Helpers.S) = struct
@@ -44,7 +44,8 @@ module Make (C : Core) (T : Lltypes_intf.S) (H : Helpers.S) = struct
     in
     (item_typ, item_size, head_size + item_size)
 
-  let get v item_typ =
+  let get v =
+    let item_typ = item_type v.typ in
     let lltyp = get_lltype_def item_typ in
 
     let value = Llvm.build_gep int_t v.value [| ci 1 |] "data" builder in

lib/mallocs.ml

@@ -10,10 +10,14 @@ module Make (Mtree : Monomorph_tree_intf.S) = struct
     | Path (m, p) -> Path (m, p @ [ index ])
     | (Single _ | Param _) as m -> Path (m, [ index ])
 
-  let m_to_list = function
+  let rec m_to_list = function
     | Malloc.No_malloc -> []
     | Single i -> [ i.mid ]
     | Param _ -> []
+    | Path (m, [ -1 ]) ->
+        (* Special case for rc/get *)
+        m_to_list m
+    | Path (m, -1 :: tl) -> m_to_list (Path (m, tl))
     | Path _ -> failwith "Internal Error: Path not supported here"
 
   type pmap = Pset.t Imap.t
@@ -96,6 +100,19 @@ module Make (Mtree : Monomorph_tree_intf.S) = struct
                 (0, true, Pset.empty) fs
             in
             (excluded, pset)
+        | Trc t ->
+            if contains_allocation t then
+              match pop_index_pset frees (-1) with
+              | Not_excl -> (false, Pset.empty)
+              | Excl -> (false, Pset.singleton [ -1 ])
+              | Followup frees ->
+                  let nexcluded, npset = is_excluded frees t in
+                  let npset =
+                    if nexcluded then Pset.empty
+                    else Pset.map (fun l -> -1 :: l) npset
+                  in
+                  (nexcluded, npset)
+            else (true, Pset.empty)
         | _ -> failwith "todo exh"
     in
     let frees =

lib/monomorph_tree.ml

@@ -1236,9 +1236,16 @@ and morph_app mk p callee args ret_typ =
     else (false, p)
   in
 
+  let fst_arg_malloc = ref None in
+
   let f p (arg, attr) =
     let ret = p.ret in
     let p, ex, var = morph_expr { p with ret = false } arg in
+
+    (match !fst_arg_malloc with
+    | Some _ -> ()
+    | None -> fst_arg_malloc := Some var.malloc);
+
     let is_moved =
       match attr with Typed_tree.Dmove -> true | Dset | Dmut | Dnorm -> false
     in
@@ -1299,8 +1306,11 @@ and morph_app mk p callee args ret_typ =
     (* array-get does not return a temporary. If its value is returned in a
        function, increase value's refcount so that it's really a temporary *)
     match callee.monomorph with
-    | Builtin ((Array_get | Fixed_array_get | Unsafe_ptr_get | Rc_get), _) ->
+    | Builtin ((Array_get | Fixed_array_get | Unsafe_ptr_get), _) ->
         (Malloc.No_malloc, p.mallocs)
+    | Builtin (Rc_get, _) ->
+        let malloc = malloc_add_index (-1) (Option.get !fst_arg_malloc) in
+        (malloc, p.mallocs)
     | _ ->
         let _, malloc, mallocs = mb_malloc None p.mallocs ret_typ in
         (malloc, mallocs)
@@ -1365,11 +1375,11 @@ and morph_var_data mk p expr typ =
   in
   let func =
     (* Since we essentially change the datatype here, we have to be sure that
-       the variant was allocated before. Usually it is, but in the case of toplevel
-       lets it might not. For instance if we have an (option t) which is matched on
-       at assignment. Then, the global value is t, but if we propagate the alloc,
-       the parent (option t) will try to initialize into the global value, which is t,
-       another type.*)
+       the variant was allocated before. Usually it is, but in the case of
+       toplevel lets it might not. For instance if we have an (option t) which
+       is matched on at assignment. Then, the global value is t, but if we
+       propagate the alloc, the parent (option t) will try to initialize into
+       the global value, which is t, another type.*)
     if p.toplvl then
       let alloc = Value (ref (request p)) in
       { func with alloc }

test/memory.t/partial_rc.smu

@@ -0,0 +1,8 @@
+fun drop_rc(str&, rc!) {
+  match rc/get(rc) {
+    v: &str = v
+  }
+}
+
+let str& = ""
+drop_rc(&str, !rc/create(!",.p"))

test/memory.t/run.t

@@ -3157,3 +3157,7 @@ Mutable variables in upward closures
   2
   3
   4
+
+Partially free rc
+  $ schmu partial_rc.smu
+  $ valgrind -q --leak-check=yes --show-reachable=yes ./partial_rc