From 568bd1c269cd61bd2099c9ffcdcf780de53dfab3 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Wed, 22 May 2019 11:44:09 -0700
Subject: [PATCH] fix(library/init/lean/compiler/ir/borrow): do not perform
 borrow inference for `@[export]` declarations

---
 library/init/lean/compiler/export.lean    |  3 +++
 library/init/lean/compiler/ir/borrow.lean | 30 ++++++++++++++++-------
 2 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/library/init/lean/compiler/export.lean b/library/init/lean/compiler/export.lean
index 2f0a1db317..b7d4a66d77 100644
--- a/library/init/lean/compiler/export.lean
+++ b/library/init/lean/compiler/export.lean
@@ -9,4 +9,7 @@ import init.lean.environment
 namespace Lean
 @[extern "lean_get_export_name_for"]
 constant getExportNameFor (env : @& Environment) (n : @& Name) : Option Name := default _
+
+def isExport (env : Environment) (n : Name) : Bool :=
+(getExportNameFor env n).isSome
 end Lean
diff --git a/library/init/lean/compiler/ir/borrow.lean b/library/init/lean/compiler/ir/borrow.lean
index 5f1a2e354a..73e8068d4d 100644
--- a/library/init/lean/compiler/ir/borrow.lean
+++ b/library/init/lean/compiler/ir/borrow.lean
@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 prelude
+import init.lean.compiler.export
 import init.lean.compiler.ir.compilerm
 import init.lean.compiler.ir.normids
 
@@ -51,26 +52,37 @@ namespace InitParamMap
 def initBorrow (ps : Array Param) : Array Param :=
 ps.hmap $ λ p, { borrow := p.ty.isObj, .. p }
 
-partial def visitFnBody : FnBody → ReaderT FunId (State ParamMap) Unit
+/- We do perform borrow inference for constants marked as `export`.
+   Reason: we current write wrappers in C++ for using exported functions.
+   These wrappers use smart pointers such as `object_ref`.
+   When writing a new wrapper we need to know whether an argument is a borrow
+   inference or not.
+
+   We can revise this decision when we implement code for generating
+   the wrappers automatically. -/
+def initBorrowIfNotExported (exported : Bool) (ps : Array Param) : Array Param :=
+if exported then ps else initBorrow ps
+
+partial def visitFnBody (exported : Bool) (fnid : FunId) : FnBody → State ParamMap Unit
 | (FnBody.jdecl j xs v b) := do
-  fnid ← read,
-  modify $ λ m, m.insert (Key.jp fnid j) (initBorrow xs),
+  modify $ λ m, m.insert (Key.jp fnid j) (initBorrowIfNotExported exported xs),
   visitFnBody b
 | e :=
   unless (e.isTerminal) $ do
     let (instr, b) := e.split,
     visitFnBody b
 
-def visitDecls (decls : Array Decl) : State ParamMap Unit :=
+def visitDecls (env : Environment) (decls : Array Decl) : State ParamMap Unit :=
 decls.mfor $ λ decl, match decl with
   | Decl.fdecl f xs _ b := do
-    modify $ λ m, m.insert (Key.decl f) (initBorrow xs),
-    visitFnBody b f
+    let exported := isExport env f,
+    modify $ λ m, m.insert (Key.decl f) (initBorrowIfNotExported exported xs),
+    visitFnBody exported f b
   | _ := pure ()
 end InitParamMap
 
-def mkInitParamMap (decls : Array Decl) : ParamMap :=
-(InitParamMap.visitDecls decls *> get).run' {}
+def mkInitParamMap (env : Environment) (decls : Array Decl) : ParamMap :=
+(InitParamMap.visitDecls env decls *> get).run' {}
 
 namespace ApplyParamMap
 
@@ -276,7 +288,7 @@ s ← get,
 pure s.map
 
 def infer (env : Environment) (decls : Array Decl) : ParamMap :=
-(collectDecls decls { env := env }).run' { map := mkInitParamMap decls }
+(collectDecls decls { env := env }).run' { map := mkInitParamMap env decls }
 
 end Borrow