From 664b72deabed439e54e24a7d4bce0b4247638429 Mon Sep 17 00:00:00 2001
From: Aymeric Fromherz <aymeric.fromherz@inria.fr>
Date: Tue, 4 Feb 2025 11:24:46 +0100
Subject: [PATCH] Support assumed functions in Rust mutability inference

---
 lib/AstToMiniRust.ml    | 9 +++++++--
 lib/OptimizeMiniRust.ml | 5 +++++
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/lib/AstToMiniRust.ml b/lib/AstToMiniRust.ml
index 4edfed00..0a9ab653 100644
--- a/lib/AstToMiniRust.ml
+++ b/lib/AstToMiniRust.ml
@@ -1400,8 +1400,13 @@ let translate_decl env (d: Ast.decl): MiniRust.decl option =
       let meta = translate_meta flags in
       Some (MiniRust.Constant { name; typ; body; meta })
 
-  | DExternal _ ->
-      None
+  | DExternal (_, _, _, _, lid, _, _) ->
+      let name, parameters, return_type =
+        match lookup_decl env lid with
+        | name, Function (_, parameters, return_type) -> name, parameters, return_type
+        | _ -> failwith " impossible"
+      in
+      Some (MiniRust.Assumed { name; parameters; return_type })
 
   | DType (lid, flags, _, _, decl) ->
       let name = lookup_type env lid in
diff --git a/lib/OptimizeMiniRust.ml b/lib/OptimizeMiniRust.ml
index 64b3a030..09ba2807 100644
--- a/lib/OptimizeMiniRust.ml
+++ b/lib/OptimizeMiniRust.ml
@@ -669,6 +669,8 @@ let infer_function (env: env) valuation (d: decl): decl =
          the traversal does not add or remove any bindings, but only increases the
          mutability, we can do a direct replacement instead of a more complex merge *)
       Function { f with body; parameters }
+  (* Assumed functions already have their mutability specified, we skip them *)
+  | Assumed _ -> d
   | _ ->
       assert false
 
@@ -1021,6 +1023,8 @@ let infer_mut_borrows files =
         List.filter_map (function
           | Function { parameters; name; _ } ->
               Some (name, List.map (fun (p: MiniRust.binding) -> p.typ) parameters)
+          | Assumed { name; parameters; _ } ->
+              Some (name, parameters)
           | _ ->
               None
         ) decls) files))
@@ -1044,6 +1048,7 @@ let infer_mut_borrows files =
     else
       match infer_function env valuation (NameMap.find name definitions) with
       | Function { parameters; _ } -> distill (List.map (fun (b: MiniRust.binding) -> b.typ) parameters)
+      | Assumed { parameters; _ } -> distill parameters
       | _ -> failwith "impossible"
   in