From 9f64f53fefcb08cbcd85af5a27d1ad08c5450619 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Henrik=20B=C3=B6ving?= Date: Fri, 13 Feb 2026 10:56:50 +0100 Subject: [PATCH] refactor: port Boxing from IR to LCNF (#12458) This PR ports the IR pass for box/unbox insertion to LCNF. --- src/Lean/Compiler/IR.lean | 4 +- src/Lean/Compiler/IR/AddExtern.lean | 32 +- src/Lean/Compiler/IR/Boxing.lean | 341 ------------------- src/Lean/Compiler/IR/CompilerM.lean | 13 +- src/Lean/Compiler/IR/EmitC.lean | 7 +- src/Lean/Compiler/IR/EmitLLVM.lean | 7 +- src/Lean/Compiler/IR/ToIR.lean | 6 + src/Lean/Compiler/LCNF/AlphaEqv.lean | 2 + src/Lean/Compiler/LCNF/Basic.lean | 52 ++- src/Lean/Compiler/LCNF/CompilerM.lean | 8 + src/Lean/Compiler/LCNF/DependsOn.lean | 2 +- src/Lean/Compiler/LCNF/ElimDead.lean | 10 +- src/Lean/Compiler/LCNF/ExplicitBoxing.lean | 378 +++++++++++++++++++++ src/Lean/Compiler/LCNF/FVarUtil.lean | 4 +- src/Lean/Compiler/LCNF/Internalize.lean | 9 +- src/Lean/Compiler/LCNF/Passes.lean | 2 + src/Lean/Compiler/LCNF/PrettyPrinter.lean | 2 + tests/lean/externBoxing.lean.expected.out | 2 +- tests/lean/run/mangling.lean | 5 +- 19 files changed, 495 insertions(+), 391 deletions(-) delete mode 100644 src/Lean/Compiler/IR/Boxing.lean create mode 100644 src/Lean/Compiler/LCNF/ExplicitBoxing.lean diff --git a/src/Lean/Compiler/IR.lean b/src/Lean/Compiler/IR.lean index 6d6a69af32..df4250638e 100644 --- a/src/Lean/Compiler/IR.lean +++ b/src/Lean/Compiler/IR.lean @@ -13,7 +13,6 @@ public import Lean.Compiler.IR.CompilerM public import Lean.Compiler.IR.PushProj public import Lean.Compiler.IR.NormIds public import Lean.Compiler.IR.Checker -public import Lean.Compiler.IR.Boxing public import Lean.Compiler.IR.RC public import Lean.Compiler.IR.ExpandResetReuse public import Lean.Compiler.IR.UnboxResult @@ -24,6 +23,7 @@ public import Lean.Compiler.IR.ToIRType public import Lean.Compiler.IR.Meta public import Lean.Compiler.IR.Toposort public import Lean.Compiler.IR.SimpleGroundExpr +public import Lean.Compiler.IR.ElimDeadVars -- The following imports are not required by the compiler. They are here to ensure that there -- are no orphaned modules. @@ -38,8 +38,6 @@ def compile (decls : Array Decl) : CompilerM (Array Decl) := do logDecls `init decls checkDecls decls let mut decls := decls - decls ← explicitBoxing decls - logDecls `boxing decls decls ← explicitRC decls logDecls `rc decls if Compiler.LCNF.compiler.reuse.get (← getOptions) then diff --git a/src/Lean/Compiler/IR/AddExtern.lean b/src/Lean/Compiler/IR/AddExtern.lean index 7465ca04db..8f9c4ce3e3 100644 --- a/src/Lean/Compiler/IR/AddExtern.lean +++ b/src/Lean/Compiler/IR/AddExtern.lean @@ -7,12 +7,14 @@ Authors: Cameron Zwarich module prelude -public import Lean.Compiler.IR.Boxing +import Init.While import Lean.Compiler.IR.RC +import Lean.Compiler.IR.ToIR import Lean.Compiler.LCNF.ToImpureType import Lean.Compiler.LCNF.ToImpure -import Init.While -import Lean.Compiler.LCNF.PhaseExt +import Lean.Compiler.LCNF.ExplicitBoxing +import Lean.Compiler.LCNF.Internalize +public import Lean.Compiler.ExternAttr public section @@ -23,8 +25,8 @@ def addExtern (declName : Name) (externAttrData : ExternAttrData) : CoreM Unit : if !isPrivateName declName then modifyEnv (Compiler.LCNF.setDeclPublic · declName) let monoDecl ← addMono declName - let impureDecl ← addImpure monoDecl - addIr impureDecl + let impureDecls ← addImpure monoDecl + addIr impureDecls where addMono (declName : Name) : CoreM (Compiler.LCNF.Decl .pure) := do let type ← Compiler.LCNF.getOtherDeclMonoType declName @@ -51,11 +53,11 @@ where decl.saveMono return decl - addImpure (decl : Compiler.LCNF.Decl .pure) : CoreM (Compiler.LCNF.Decl .impure) := do + addImpure (decl : Compiler.LCNF.Decl .pure) : CoreM (Array (Compiler.LCNF.Decl .impure)) := do let type ← Compiler.LCNF.lowerResultType decl.type decl.params.size let params ← decl.params.mapM fun param => return { param with type := ← Compiler.LCNF.toImpureType param.type } - let decl := { + let decl : Compiler.LCNF.Decl .impure := { name := decl.name, levelParams := decl.levelParams, value := .extern externAttrData @@ -63,15 +65,13 @@ where type, params } - decl.saveImpure - return decl - - addIr (decl : Compiler.LCNF.Decl .impure) : CoreM Unit := do - let params := decl.params.mapIdx fun idx param => - { x := ⟨idx⟩, borrow := param.borrow, ty := toIRType param.type } - let type := toIRType decl.type - let decls := #[.extern declName params type externAttrData] - let decls := ExplicitBoxing.addBoxedVersions (← Lean.getEnv) decls + Compiler.LCNF.CompilerM.run (phase := .impure) do + let decl ← decl.internalize + decl.saveImpure + Compiler.LCNF.addBoxedVersions #[decl] + + addIr (decls : Array (Compiler.LCNF.Decl .impure)) : CoreM Unit := do + let decls ← toIR decls let decls ← explicitRC decls logDecls `result decls addDecls decls diff --git a/src/Lean/Compiler/IR/Boxing.lean b/src/Lean/Compiler/IR/Boxing.lean deleted file mode 100644 index c8a1e4bd47..0000000000 --- a/src/Lean/Compiler/IR/Boxing.lean +++ /dev/null @@ -1,341 +0,0 @@ -/- -Copyright (c) 2019 Microsoft Corporation. All rights reserved. -Released under Apache 2.0 license as described in the file LICENSE. -Authors: Leonardo de Moura --/ -module - -prelude -public import Lean.Runtime -public import Lean.Compiler.ClosedTermCache -public import Lean.Compiler.IR.CompilerM -public import Lean.Compiler.IR.ElimDeadVars -public import Lean.Compiler.IR.ToIRType -public import Lean.Data.AssocList - -public section - -namespace Lean.IR.ExplicitBoxing -/-! -Add explicit boxing and unboxing instructions. -Recall that the Lean to λ_pure compiler produces code without these instructions. - -Assumptions: -- This transformation is applied before explicit RC instructions (`inc`, `dec`) are inserted. -- This transformation is applied before lower level optimizations are applied which use - `Expr.isShared`, `Expr.isTaggedPtr`, and `FnBody.set`. --/ - -abbrev N := StateM Nat - -private def N.mkFresh : N VarId := - modifyGet fun n => ({ idx := n }, n + 1) - -def requiresBoxedVersion (env : Environment) (decl : Decl) : Bool := - let ps := decl.params - (ps.size > 0 && (decl.resultType.isScalar || ps.any (fun p => p.ty.isScalar || p.borrow || p.ty.isVoid) || isExtern env decl.name)) - || ps.size > closureMaxArgs - -def mkBoxedVersionAux (decl : Decl) : N Decl := do - let ps := decl.params - let qs ← ps.mapM fun p => do let x ← N.mkFresh; pure { x, ty := p.ty.boxed, borrow := false } - let (newVDecls, xs) ← qs.size.foldM (init := (#[], #[])) fun i _ (newVDecls, xs) => do - let p := ps[i]! - let q := qs[i] - if !p.ty.isScalar then - pure (newVDecls, xs.push (.var q.x)) - else - let x ← N.mkFresh - pure (newVDecls.push (.vdecl x p.ty (.unbox q.x) default), xs.push (.var x)) - let r ← N.mkFresh - let newVDecls := newVDecls.push (.vdecl r decl.resultType (.fap decl.name xs) default) - let body ← if !decl.resultType.isScalar then - pure <| reshape newVDecls (.ret (.var r)) - else - let newR ← N.mkFresh - let newVDecls := newVDecls.push (.vdecl newR decl.resultType.boxed (.box decl.resultType r) default) - pure <| reshape newVDecls (.ret (.var newR)) - return Decl.fdecl (mkBoxedName decl.name) qs decl.resultType.boxed body decl.getInfo - -def mkBoxedVersion (decl : Decl) : Decl := - (mkBoxedVersionAux decl).run' 1 - -def addBoxedVersions (env : Environment) (decls : Array Decl) : Array Decl := - let boxedDecls := decls.foldl (init := #[]) fun newDecls decl => - if requiresBoxedVersion env decl then newDecls.push (mkBoxedVersion decl) else newDecls - decls ++ boxedDecls - -def eqvTypes (t₁ t₂ : IRType) : Bool := - (t₁.isScalar == t₂.isScalar) && (!t₁.isScalar || t₁ == t₂) - -structure BoxingContext where - f : FunId - localCtx : LocalContext := {} - resultType : IRType - decls : Array Decl - env : Environment - -structure BoxingState where - nextIdx : Index - /-- We create auxiliary declarations when boxing constant and literals. - The idea is to avoid code such as - ``` - let x1 := Uint64.inhabited; - let x2 := box x1; - ... - ``` - We currently do not cache these declarations in an environment extension, but - we use auxDeclCache to avoid creating equivalent auxiliary declarations more than once when - processing the same IR declaration. - -/ - auxDecls : Array Decl := #[] - auxDeclCache : AssocList FnBody Expr := AssocList.empty - nextAuxId : Nat := 1 - -abbrev M := ReaderT BoxingContext (StateT BoxingState Id) - -private def M.mkFresh : M VarId := do - let oldS ← getModify fun s => { s with nextIdx := s.nextIdx + 1 } - pure { idx := oldS.nextIdx } - -def getEnv : M Environment := BoxingContext.env <$> read -def getLocalContext : M LocalContext := BoxingContext.localCtx <$> read -def getResultType : M IRType := BoxingContext.resultType <$> read - -def getVarType (x : VarId) : M IRType := do - let localCtx ← getLocalContext - match localCtx.getType x with - | some t => pure t - | none => pure .tobject -- unreachable, we assume the code is well formed - -def getJPParams (j : JoinPointId) : M (Array Param) := do - let localCtx ← getLocalContext - match localCtx.getJPParams j with - | some ys => pure ys - | none => pure #[] -- unreachable, we assume the code is well formed - -def getDecl (fid : FunId) : M Decl := do - let ctx ← read - match findEnvDecl' ctx.env fid ctx.decls with - | some decl => pure decl - | none => pure default -- unreachable if well-formed - -@[inline] def withParams {α : Type} (xs : Array Param) (k : M α) : M α := - withReader (fun ctx => { ctx with localCtx := ctx.localCtx.addParams xs }) k - -@[inline] def withVDecl {α : Type} (x : VarId) (ty : IRType) (v : Expr) (k : M α) : M α := - withReader (fun ctx => { ctx with localCtx := ctx.localCtx.addLocal x ty v }) k - -@[inline] def withJDecl {α : Type} (j : JoinPointId) (xs : Array Param) (v : FnBody) (k : M α) : M α := - withReader (fun ctx => { ctx with localCtx := ctx.localCtx.addJP j xs v }) k - -/-- If `x` declaration is of the form `x := Expr.lit _` or `x := Expr.fap c #[]`, - and `x`'s type is not cheap to box (e.g., it is `UInt64), then return its value. -/ -private def isExpensiveConstantValueBoxing (x : VarId) (xType : IRType) : M (Option Expr) := - match xType with - | .uint8 | .uint16 => return none - | _ => do - let localCtx ← getLocalContext - match localCtx.getValue x with - | some val => - match val with - -- TODO: This should check whether larger literals fit into tagged values. - | .lit _ => return some val - | .fap _ args => return if args.size == 0 then some val else none - | _ => return none - | _ => return none - -/-- Auxiliary function used by castVarIfNeeded. - It is used when the expected type does not match `xType`. - If `xType` is scalar, then we need to "box" it. Otherwise, we need to "unbox" it. -/ -def mkCast (x : VarId) (xType : IRType) (expectedType : IRType) : M Expr := do - if expectedType.isScalar then - return .unbox x - else - match (← isExpensiveConstantValueBoxing x xType) with - | some v => do - let ctx ← read - let s ← get - /- Create auxiliary FnBody - ``` - let x_1 : xType := v; - let x_2 : expectedType := Expr.box xType x_1; - ret x_2 - ``` - -/ - let body : FnBody := - .vdecl { idx := 1 } xType v <| - .vdecl { idx := 2 } expectedType (.box xType { idx := 1 }) <| - .ret (.var { idx := 2 }) - match s.auxDeclCache.find? body with - | some v => pure v - | none => do - let auxName := ctx.f ++ ((`_boxed_const).appendIndexAfter s.nextAuxId) - let auxConst := .fap auxName #[] - let auxDecl := Decl.fdecl auxName #[] expectedType body {} - modify fun s => { s with - auxDecls := s.auxDecls.push auxDecl - auxDeclCache := s.auxDeclCache.cons body auxConst - nextAuxId := s.nextAuxId + 1 - } - pure auxConst - | none => return .box xType x - -@[inline] def castVarIfNeeded (x : VarId) (expected : IRType) (k : VarId → M FnBody) : M FnBody := do - let xType ← getVarType x - if eqvTypes xType expected then - k x - else - let y ← M.mkFresh - let v ← mkCast x xType expected - .vdecl y expected v <$> k y - -@[inline] def castArgIfNeeded (x : Arg) (expected : IRType) (k : Arg → M FnBody) : M FnBody := - match x with - | .var x => castVarIfNeeded x expected (fun x => k (.var x)) - | .erased => k x - -def castArgsIfNeededAux (xs : Array Arg) (typeFromIdx : Nat → IRType) : M (Array Arg × Array FnBody) := do - let mut xs' := #[] - let mut bs := #[] - let mut i := 0 - for x in xs do - let expected := typeFromIdx i - match x with - | .erased => - xs' := xs'.push x - | .var x => - let xType ← getVarType x - if eqvTypes xType expected then - xs' := xs'.push (.var x) - else - let y ← M.mkFresh - let v ← mkCast x xType expected - let b := .vdecl y expected v .nil - xs' := xs'.push (.var y) - bs := bs.push b - i := i + 1 - return (xs', bs) - -@[inline] def castArgsIfNeeded (xs : Array Arg) (ps : Array Param) (k : Array Arg → M FnBody) : M FnBody := do - let (ys, bs) ← castArgsIfNeededAux xs fun i => ps[i]!.ty - let b ← k ys - pure (reshape bs b) - -@[inline] def boxArgsIfNeeded (xs : Array Arg) (k : Array Arg → M FnBody) : M FnBody := do - let (ys, bs) ← castArgsIfNeededAux xs (fun _ => .tobject) - let b ← k ys - pure (reshape bs b) - -def unboxResultIfNeeded (x : VarId) (ty : IRType) (e : Expr) (b : FnBody) : M FnBody := do - if ty.isScalar then - let y ← M.mkFresh - return .vdecl y .tobject e (.vdecl x ty (.unbox y) b) - else - return .vdecl x ty e b - -def castResultIfNeeded (x : VarId) (ty : IRType) (e : Expr) (eType : IRType) (b : FnBody) : M FnBody := do - if eqvTypes ty eType then - return .vdecl x ty e b - else - let y ← M.mkFresh - let boxedTy := ty.boxed - let v ← mkCast y boxedTy ty - return .vdecl y boxedTy e (.vdecl x ty v b) - -def visitVDeclExpr (x : VarId) (ty : IRType) (e : Expr) (b : FnBody) : M FnBody := - match e with - | .ctor c ys => - if c.isScalar && ty.isScalar then - return .vdecl x ty (.lit (.num c.cidx)) b - else - boxArgsIfNeeded ys fun ys => return .vdecl x ty (.ctor c ys) b - | .reuse w c u ys => - boxArgsIfNeeded ys fun ys => return .vdecl x ty (.reuse w c u ys) b - | .fap f ys => do - let decl ← getDecl f - castArgsIfNeeded ys decl.params fun ys => - castResultIfNeeded x ty (.fap f ys) decl.resultType b - | .pap f ys => do - let env ← getEnv - let decl ← getDecl f - let f := if requiresBoxedVersion env decl then mkBoxedName f else f - boxArgsIfNeeded ys fun ys => return .vdecl x ty (.pap f ys) b - | .ap f ys => - boxArgsIfNeeded ys fun ys => - unboxResultIfNeeded x ty (.ap f ys) b - | _ => - return .vdecl x ty e b - -/-- -Up to this point the type system of IR is quite loose so we can for example encounter situations -such as -``` -let y : obj := f x -``` -where `f : obj -> uint8`. It is the job of the boxing pass to enforce a stricter obj/scalar -separation and as such it needs to correct situations like this. --/ -def tryCorrectVDeclType (ty : IRType) (e : Expr) : M IRType := - match e with - | .fap f _ => do - let decl ← getDecl f - return decl.resultType - | .pap .. => return .object - | .uproj .. => return .usize - | .ctor .. | .reuse .. | .ap .. | .lit .. | .sproj .. | .proj .. | .reset .. => - return ty - | .unbox .. | .box .. | .isShared .. => unreachable! - -partial def visitFnBody : FnBody → M FnBody - | .vdecl x t v b => do - let t ← tryCorrectVDeclType t v - let b ← withVDecl x t v (visitFnBody b) - visitVDeclExpr x t v b - | .jdecl j xs v b => do - let v ← withParams xs (visitFnBody v) - let b ← withJDecl j xs v (visitFnBody b) - return .jdecl j xs v b - | .uset x i y b => do - let b ← visitFnBody b - castVarIfNeeded y IRType.usize fun y => - return .uset x i y b - | .sset x i o y ty b => do - let b ← visitFnBody b - castVarIfNeeded y ty fun y => - return .sset x i o y ty b - | .case tid x xType alts => do - let alts ← alts.mapM fun alt => alt.modifyBodyM visitFnBody - castVarIfNeeded x xType fun x => do - return .case tid x xType alts - | .ret x => do - let expected ← getResultType - castArgIfNeeded x expected (fun x => return .ret x) - | .jmp j ys => do - let ps ← getJPParams j - castArgsIfNeeded ys ps fun ys => return .jmp j ys - | other => - pure other - -def run (env : Environment) (decls : Array Decl) : Array Decl := - let decls := decls.foldl (init := #[]) fun newDecls decl => - match decl with - | .fdecl f xs resultType b _ => - let nextIdx := decl.maxIndex + 1 - let (b, s) := withParams xs (visitFnBody b) { f, resultType, decls, env } |>.run { nextIdx } - let newDecls := newDecls ++ s.auxDecls - let newDecl := decl.updateBody! b - let newDecl := newDecl.elimDead - newDecls.push newDecl - | d => newDecls.push d - addBoxedVersions env decls - -end ExplicitBoxing - -def explicitBoxing (decls : Array Decl) : CompilerM (Array Decl) := do - let env ← getEnv - return ExplicitBoxing.run env decls - -builtin_initialize registerTraceClass `compiler.ir.boxing (inherited := true) - -end Lean.IR diff --git a/src/Lean/Compiler/IR/CompilerM.lean b/src/Lean/Compiler/IR/CompilerM.lean index 2551d83f69..8b26631019 100644 --- a/src/Lean/Compiler/IR/CompilerM.lean +++ b/src/Lean/Compiler/IR/CompilerM.lean @@ -11,6 +11,7 @@ public import Lean.Compiler.ExportAttr public import Lean.Compiler.LCNF.PublicDeclsExt import Lean.Compiler.InitAttr import Init.Data.Format.Macro +import Lean.Compiler.LCNF.ExplicitBoxing public section @@ -140,20 +141,10 @@ def findEnvDecl (env : Environment) (declName : Name) (includeServer := false): private def findInterpDecl (env : Environment) (declName : Name) : Option Decl := findEnvDecl (includeServer := true) env declName -namespace ExplicitBoxing - -def mkBoxedName (n : Name) : Name := - Name.mkStr n "_boxed" - -def isBoxedName (name : Name) : Bool := - name matches .str _ "_boxed" - -end ExplicitBoxing - /-- Like ``findInterpDecl env (declName ++ `_boxed)`` but with optimized negative lookup. -/ @[export lean_ir_find_env_decl_boxed] private def findInterpDeclBoxed (env : Environment) (declName : Name) : Option Decl := - let boxed := ExplicitBoxing.mkBoxedName declName + let boxed := Compiler.LCNF.mkBoxedName declName -- Important: get module index of base name, not boxed version. Usually the interpreter never -- does negative lookups except in the case of `call_boxed` which must check whether a boxed -- version exists. If `declName` exists as an imported declaration but `declName'` doesn't, the diff --git a/src/Lean/Compiler/IR/EmitC.lean b/src/Lean/Compiler/IR/EmitC.lean index 07f1449052..4cbc858e1d 100644 --- a/src/Lean/Compiler/IR/EmitC.lean +++ b/src/Lean/Compiler/IR/EmitC.lean @@ -10,16 +10,19 @@ public import Lean.Compiler.NameMangling public import Lean.Compiler.IR.EmitUtil public import Lean.Compiler.IR.NormIds public import Lean.Compiler.IR.SimpCase -public import Lean.Compiler.IR.Boxing public import Lean.Compiler.ModPkgExt +import Lean.Compiler.LCNF.ExplicitBoxing +import Lean.Compiler.ClosedTermCache import Lean.Compiler.IR.SimpleGroundExpr import Init.Omega import Init.While +import Init.Data.Range.Polymorphic.Iterators +import Lean.Runtime public section namespace Lean.IR.EmitC -open ExplicitBoxing (isBoxedName) +open Lean.Compiler.LCNF (isBoxedName) def leanMainFn := "_lean_main" diff --git a/src/Lean/Compiler/IR/EmitLLVM.lean b/src/Lean/Compiler/IR/EmitLLVM.lean index a8c3c1d66a..59d6b9dbfc 100644 --- a/src/Lean/Compiler/IR/EmitLLVM.lean +++ b/src/Lean/Compiler/IR/EmitLLVM.lean @@ -10,13 +10,16 @@ public import Lean.Compiler.NameMangling public import Lean.Compiler.IR.EmitUtil public import Lean.Compiler.IR.NormIds public import Lean.Compiler.IR.SimpCase -public import Lean.Compiler.IR.Boxing public import Lean.Compiler.IR.LLVMBindings +import Lean.Compiler.LCNF.ExplicitBoxing import Lean.Compiler.ModPkgExt +import Lean.Runtime +import Lean.Compiler.ClosedTermCache +import Init.Data.Range.Polymorphic.Iterators public section -open Lean.IR.ExplicitBoxing (isBoxedName) +open Lean.Compiler.LCNF (isBoxedName) namespace Lean.IR /- diff --git a/src/Lean/Compiler/IR/ToIR.lean b/src/Lean/Compiler/IR/ToIR.lean index 8f91170daa..4140624080 100644 --- a/src/Lean/Compiler/IR/ToIR.lean +++ b/src/Lean/Compiler/IR/ToIR.lean @@ -143,6 +143,12 @@ partial def lowerLet (decl : LCNF.LetDecl .impure) (k : LCNF.Code .impure) : M F withGetFVarValue var fun var => do let irArgs ← args.mapM lowerArg continueLet (.reuse var (lowerCtorInfo i) updateHeader irArgs) + | .box ty var => + withGetFVarValue var fun var => do + continueLet (.box (toIRType ty) var) + | .unbox var => + withGetFVarValue var fun var => do + continueLet (.unbox var) | .erased => mkErased () where mkErased (_ : Unit) : M FnBody := do diff --git a/src/Lean/Compiler/LCNF/AlphaEqv.lean b/src/Lean/Compiler/LCNF/AlphaEqv.lean index 8e6c14bce3..e1571dbcda 100644 --- a/src/Lean/Compiler/LCNF/AlphaEqv.lean +++ b/src/Lean/Compiler/LCNF/AlphaEqv.lean @@ -73,6 +73,8 @@ def eqvLetValue (e₁ e₂ : LetValue pu) : EqvM Bool := do | .reset n₁ v₁ _, .reset n₂ v₂ _ => pure (n₁ == n₂) <&&> eqvFVar v₁ v₂ | .reuse v₁ i₁ u₁ as₁ _, .reuse v₂ i₂ u₂ as₂ _ => pure (i₁ == i₂ && u₁ == u₂) <&&> eqvFVar v₁ v₂ <&&> eqvArgs as₁ as₂ + | .box ty₁ v₁ _, .box ty₂ v₂ _ => eqvType ty₁ ty₂ <&&> eqvFVar v₁ v₂ + | .unbox v₁ _, .unbox v₂ _ => eqvFVar v₁ v₂ | _, _ => return false @[inline] def withFVar (fvarId₁ fvarId₂ : FVarId) (x : EqvM α) : EqvM α := diff --git a/src/Lean/Compiler/LCNF/Basic.lean b/src/Lean/Compiler/LCNF/Basic.lean index 88662e25ef..5ca459c3e8 100644 --- a/src/Lean/Compiler/LCNF/Basic.lean +++ b/src/Lean/Compiler/LCNF/Basic.lean @@ -345,6 +345,14 @@ inductive LetValue (pu : Purity) where /-- `reuse x in ctor_i ys` instruction in the paper. `updateHeader` is set if the tag in the new ctor differs from the one in the old ctor and thus needs to be updated. -/ | reuse (var : FVarId) (i : CtorInfo) (updateHeader : Bool) (args : Array (Arg pu)) (h : pu = .impure := by purity_tac) + /-- + Given a scalar type `ty` and a value `fvarId : ty`, this operation returns a value of type + `tobject`. For small scalar values the result is a tagged pointer, and no memory allocation is + performed. + -/ + | box (ty : Expr) (fvarId : FVarId) (h : pu = .impure := by purity_tac) + /-- Given `fvarId : [t]object`, obtain the underlying scalar value. -/ + | unbox (fvarId : FVarId) (h : pu = .impure := by purity_tac) deriving Inhabited, BEq, Hashable def Arg.toLetValue (arg : Arg pu) : LetValue pu := @@ -396,6 +404,36 @@ private unsafe def LetValue.updateReuseImp (e : LetValue pu) (var' : FVarId) (i' @[implemented_by LetValue.updateReuseImp] opaque LetValue.updateReuse! (e : LetValue pu) (var' : FVarId) (i' : CtorInfo) (updateHeader' : Bool) (args' : Array (Arg pu)) : LetValue pu +private unsafe def LetValue.updateFapImp (e : LetValue pu) (declName' : Name) (args' : Array (Arg pu)) : LetValue pu := + match e with + | .fap declName args _ => if declName == declName' && ptrEq args args' then e else .fap declName' args' + | _ => unreachable! + +@[implemented_by LetValue.updateFapImp] opaque LetValue.updateFap! (e : LetValue pu) (declName' : Name) (args' : Array (Arg pu)) : LetValue pu + +private unsafe def LetValue.updatePapImp (e : LetValue pu) (declName' : Name) (args' : Array (Arg pu)) : LetValue pu := + match e with + | .pap declName args _ => if declName == declName' && ptrEq args args' then e else .pap declName' args' + | _ => unreachable! + +@[implemented_by LetValue.updatePapImp] opaque LetValue.updatePap! (e : LetValue pu) (declName' : Name) (args' : Array (Arg pu)) : LetValue pu + +private unsafe def LetValue.updateBoxImp (e : LetValue pu) (ty' : Expr) (fvarId' : FVarId) : LetValue pu := + match e with + | .box ty fvarId _ => if ptrEq ty ty' && fvarId == fvarId' then e else .box ty' fvarId' + | _ => unreachable! + +@[implemented_by LetValue.updateBoxImp] opaque LetValue.updateBox! (e : LetValue pu) (ty' : Expr) (fvarId' : FVarId) : LetValue pu + +private unsafe def LetValue.updateUnboxImp (e : LetValue pu) (fvarId' : FVarId) : LetValue pu := + match e with + | .unbox fvarId _ => if fvarId == fvarId' then e else .unbox fvarId' + | _ => unreachable! + +@[implemented_by LetValue.updateUnboxImp] opaque LetValue.updateUnbox! (e : LetValue pu) (fvarId' : FVarId) : LetValue pu + + + private unsafe def LetValue.updateArgsImp (e : LetValue pu) (args' : Array (Arg pu)) : LetValue pu := match e with | .const declName us args h => if ptrEq args args' then e else .const declName us args' @@ -425,6 +463,8 @@ def LetValue.toExpr (e : LetValue pu) : Expr := | .reuse var i updateHeader args _ => mkAppN (.const `reuse []) <| #[.fvar var, .const i.name [], ToExpr.toExpr updateHeader] ++ (args.map Arg.toExpr) + | .box ty var _ => mkApp2 (.const `box []) ty (.fvar var) + | .unbox var _ => mkApp (.const `unbox []) (.fvar var) structure LetDecl (pu : Purity) where fvarId : FVarId @@ -719,17 +759,17 @@ private unsafe def updateAltImp (alt : Alt pu) (ps' : Array (Param pu)) (k' : Co (offset' : Nat) (y' : FVarId) (ty' : Expr) (k' : Code pu) : Code pu @[inline] private unsafe def updateUsetImp (c : Code pu) (fvarId' : FVarId) - (offset' : Nat) (y' : FVarId) (k' : Code pu) : Code pu := + (i' : Nat) (y' : FVarId) (k' : Code pu) : Code pu := match c with - | .sset fvarId i offset y ty k _ => - if ptrEq fvarId fvarId' && offset == offset' && ptrEq y y' && ptrEq k k' then + | .uset fvarId i y k _ => + if ptrEq fvarId fvarId' && i == i' && ptrEq y y' && ptrEq k k' then c else - .uset fvarId' offset' y' k' + .uset fvarId' i' y' k' | _ => unreachable! @[implemented_by updateUsetImp] opaque Code.updateUset! (c : Code pu) (fvarId' : FVarId) - (offset' : Nat) (y' : FVarId) (k' : Code pu) : Code pu + (i' : Nat) (y' : FVarId) (k' : Code pu) : Code pu private unsafe def updateParamCoreImp (p : Param pu) (type : Expr) : Param pu := if ptrEq type p.type then @@ -1096,7 +1136,7 @@ private def collectLetValue (e : LetValue pu) (s : FVarIdHashSet) : FVarIdHashSe | .fvar fvarId args => collectArgs args <| s.insert fvarId | .const _ _ args _ | .pap _ args _ | .fap _ args _ | .ctor _ args _ => collectArgs args s | .proj _ _ fvarId _ | .sproj _ _ fvarId _ | .uproj _ fvarId _ | .oproj _ fvarId _ - | .reset _ fvarId _ => s.insert fvarId + | .reset _ fvarId _ | .box _ fvarId _ | .unbox fvarId _ => s.insert fvarId | .lit .. | .erased => s | .reuse fvarId _ _ args _ => collectArgs args <| s.insert fvarId diff --git a/src/Lean/Compiler/LCNF/CompilerM.lean b/src/Lean/Compiler/LCNF/CompilerM.lean index a7ea393682..3615c758c8 100644 --- a/src/Lean/Compiler/LCNF/CompilerM.lean +++ b/src/Lean/Compiler/LCNF/CompilerM.lean @@ -292,6 +292,14 @@ private partial def normLetValueImp (s : FVarSubst pu) (e : LetValue pu) (transl match normFVarImp s fvarId translator with | .fvar fvarId' => e.updateReuse! fvarId' info updateHeader (normArgsImp s args translator) | .erased => .erased + | .box ty fvarId _ => + match normFVarImp s fvarId translator with + | .fvar fvarId' => e.updateBox! ty fvarId' + | .erased => .erased + | .unbox fvarId _ => + match normFVarImp s fvarId translator with + | .fvar fvarId' => e.updateUnbox! fvarId' + | .erased => .erased /-- Interface for monads that have a free substitutions. diff --git a/src/Lean/Compiler/LCNF/DependsOn.lean b/src/Lean/Compiler/LCNF/DependsOn.lean index 6d3b0c1635..9875fb3bdf 100644 --- a/src/Lean/Compiler/LCNF/DependsOn.lean +++ b/src/Lean/Compiler/LCNF/DependsOn.lean @@ -31,7 +31,7 @@ private def letValueDepOn (e : LetValue pu) : M Bool := match e with | .erased | .lit .. => return false | .proj _ _ fvarId _ | .oproj _ fvarId _ | .uproj _ fvarId _ | .sproj _ _ fvarId _ - | .reset _ fvarId _ => fvarDepOn fvarId + | .reset _ fvarId _ | .box _ fvarId _ | .unbox fvarId _ => fvarDepOn fvarId | .fvar fvarId args | .reuse fvarId _ _ args _ => fvarDepOn fvarId <||> args.anyM argDepOn | .const _ _ args _ | .ctor _ args _ | .fap _ args _ | .pap _ args _ => args.anyM argDepOn diff --git a/src/Lean/Compiler/LCNF/ElimDead.lean b/src/Lean/Compiler/LCNF/ElimDead.lean index 9063d7a593..76403b2141 100644 --- a/src/Lean/Compiler/LCNF/ElimDead.lean +++ b/src/Lean/Compiler/LCNF/ElimDead.lean @@ -35,7 +35,7 @@ def collectLocalDeclsLetValue (s : UsedLocalDecls) (e : LetValue pu) : UsedLocal match e with | .erased | .lit .. => s | .proj _ _ fvarId _ | .reset _ fvarId _ | .sproj _ _ fvarId _ | .uproj _ fvarId _ - | .oproj _ fvarId _ => s.insert fvarId + | .oproj _ fvarId _ | .box _ fvarId _ | .unbox fvarId _ => s.insert fvarId | .const _ _ args _ => collectLocalDeclsArgs s args | .fvar fvarId args | .reuse fvarId _ _ args _ => collectLocalDeclsArgs (s.insert fvarId) args | .fap _ args _ | .pap _ args _ | .ctor _ args _ => collectLocalDeclsArgs s args @@ -56,9 +56,9 @@ def LetValue.safeToElim (val : LetValue pu) : Bool := | .pure => true | .impure => match val with + -- TODO | .isShared .. | .ctor .. | .reset .. | .reuse .. | .oproj .. | .uproj .. | .sproj .. | .lit .. | .pap .. - -- TODO | .box .. | .unbox .. | .isShared .. - | .erased .. => true + | .box .. | .unbox .. | .erased .. => true -- 0-ary full applications are considered constants | .fap _ args => args.isEmpty | .fvar .. => false @@ -105,11 +105,11 @@ partial def Code.elimDead (code : Code pu) : M (Code pu) := do end -def Decl.elimDead (decl : Decl pu) : CompilerM (Decl pu) := do +public def Decl.elimDeadVars (decl : Decl pu) : CompilerM (Decl pu) := do return { decl with value := (← decl.value.mapCodeM fun code => code.elimDead.run' {}) } public def elimDeadVars (phase : Phase) (occurrence : Nat) : Pass := - Pass.mkPerDeclaration `elimDeadVars phase Decl.elimDead occurrence + Pass.mkPerDeclaration `elimDeadVars phase Decl.elimDeadVars occurrence builtin_initialize registerTraceClass `Compiler.elimDeadVars (inherited := true) diff --git a/src/Lean/Compiler/LCNF/ExplicitBoxing.lean b/src/Lean/Compiler/LCNF/ExplicitBoxing.lean new file mode 100644 index 0000000000..a4081b197d --- /dev/null +++ b/src/Lean/Compiler/LCNF/ExplicitBoxing.lean @@ -0,0 +1,378 @@ +/- +Copyright (c) 2026 Lean FRO, LLC. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Henrik Böving +-/ +module + +prelude +public import Lean.Compiler.LCNF.CompilerM +public import Lean.Compiler.LCNF.PassManager +import Lean.Compiler.LCNF.ElimDead +import Lean.Compiler.LCNF.PhaseExt +import Lean.Compiler.LCNF.AuxDeclCache +import Lean.Runtime + +/-! +This pass is responsible for inserting `box` and `unbox` instructions and generally attempts to make +the IR actually type correct. After this pass is the first time where we can generally assume type +information to actually be correct in the entirety of LCNF. Furthermore, it also generates `boxed` +versions of functions if required. They take all arguments as [t]object/tagged and return a +[t]object/tagged. These functions are used by the interpreter and when allocating closures. + +This pass does not support: `isShared`, `inc`, `dec`, `set`, `setTag` and `del`. +-/ + +namespace Lean.Compiler.LCNF + +open ImpureType + +public def mkBoxedName (n : Name) : Name := + Name.mkStr n "_boxed" + +public def isBoxedName (name : Name) : Bool := + name matches .str _ "_boxed" + +def requiresBoxedVersion (sig : Signature .impure) : CompilerM Bool := do + let ps := sig.params + let env ← getEnv + return (ps.size > 0 + && (sig.type.isScalar + || ps.any (fun p => p.type.isScalar || p.borrow || p.type.isVoid) + || isExtern env sig.name)) + || ps.size > closureMaxArgs + +/-- +For a given signature we generate the `boxed` version by: +- declaring all parameters as the boxed variant of the current parameters +- inserting unbox instructions as required +- invoking the function +- boxing the result if required before returning it +-/ +def mkBoxedVersion (sig : Signature .impure) : CompilerM (Decl .impure) := do + let newParams ← sig.params.mapM fun p => mkParam p.binderName p.type.boxed false + let mut body := #[] + let mut args := Array.emptyWithCapacity sig.params.size + for oldParam in sig.params, newParam in newParams do + if !oldParam.type.isScalar then + args := args.push <| .fvar newParam.fvarId + else + let decl ← mkLetDecl (oldParam.binderName.str "boxed") oldParam.type (.unbox newParam.fvarId) + body := body.push <| .let decl + args := args.push <| .fvar decl.fvarId + let appDecl ← mkLetDecl `res sig.type (.fap sig.name args) + body := body.push <| .let appDecl + let value ← + if !sig.type.isScalar then + pure <| attachCodeDecls body (.return appDecl.fvarId) + else + let decl ← mkLetDecl `r sig.type.boxed (.box sig.type appDecl.fvarId) + body := body.push <| .let decl + pure <| attachCodeDecls body (.return decl.fvarId) + let decl := { + name := mkBoxedName sig.name + levelParams := [] + type := sig.type.boxed + params := newParams + value := .code value + inlineAttr? := none + } + decl.saveImpure + return decl + +/-- +For all declarations in `decls` add their `_boxed` version if required. +-/ +public def addBoxedVersions (decls : Array (Decl .impure)) : CompilerM (Array (Decl .impure)) := do + let boxedDecls ← decls.filterMapM fun decl => do + if ← requiresBoxedVersion decl.toSignature then + let boxed ← mkBoxedVersion decl.toSignature + return some boxed + else + return none + return decls ++ boxedDecls + +structure Ctx where + /-- + The name of the declaration we are currently operating on. + -/ + currDecl : Name + /-- + The result type of the declaration we are currently operating on. + -/ + currDeclResultType : Expr + /-- + The SCC of declarations we are operating on. + -/ + decls : Array (Decl .impure) + +structure State where + /-- + When boxing constants and literals we generate auxiliary declarations. + This is to avoid code like: + ``` + let x1 := UInt64.inhabited; + let x2 := box x1; + ... + ``` + We cache these declarations on a per-module level but not globally through `cacheAuxDecl`. + -/ + auxDecls : Array (Decl .impure) := #[] + /-- + Counter for generating unique auxDecl names. + -/ + nextAuxIdx : Nat := 1 + +abbrev BoxM := ReaderT Ctx StateRefT State CompilerM + +@[inline] +def getResultType : BoxM Expr := return (← read).currDeclResultType + +def typesEqvForBoxing (t₁ t₂ : Expr) : Bool := + (t₁.isScalar == t₂.isScalar) && (!t₁.isScalar || t₁ == t₂) + +/-- +If `x` declaration is of the form `x := .lit _` or `x := .fap c #[]`, +and `x`'s type is not cheap to box (e.g., it is `UInt64), then return its value. +-/ +def isExpensiveConstantValueBoxing (x : FVarId) (xType : Expr) : + BoxM (Option (LetValue .impure)) := + match xType with + | uint8 | uint16 => return none + | _ => do + let some val ← findLetValue? x | return none + match val with + | .lit _ => return some val + | .fap _ args => return if args.size == 0 then some val else none + | _ => return none + +/-- +Auxiliary function used by `castVarIfNeeded`. +It is used when the expected type does not match `fvarIdType`. +If `fvarIdType` is scalar, then we need to box it. Otherwise, we need to unbox it. +-/ +def mkCast (fvarId : FVarId) (fvarIdType : Expr) (expectedType : Expr) : + BoxM (LetValue .impure) := do + if expectedType.isScalar then + return .unbox fvarId + else + match ← isExpensiveConstantValueBoxing fvarId fvarIdType with + | none => return .box fvarIdType fvarId + | some v => + /- + v is guaranteed to be closed so we can generate the following: + let _x.1 : fvarIdType := v; + let _x.2 : expectedType := box fvarIdType _x.1; + return _x.2 + -/ + let x1 ← mkLetDecl .anonymous fvarIdType v + let x2 ← mkLetDecl .anonymous expectedType (.box fvarIdType x1.fvarId) + let body : Code .impure := .let x1 <| .let x2 <| .return x2.fvarId + let auxDecl : Decl .impure := { + name := (← read).currDecl ++ ((`_boxed_const).appendIndexAfter (← get).nextAuxIdx) + levelParams := [] + type := expectedType + params := #[] + value := .code body + inlineAttr? := none + } + match ← cacheAuxDecl auxDecl with + | .alreadyCached auxName => + auxDecl.erase + let auxConst := .fap auxName #[] + return auxConst + | .new => + modify fun s => { s with + auxDecls := s.auxDecls.push auxDecl + nextAuxIdx := s.nextAuxIdx + 1 + } + auxDecl.saveImpure + let auxConst := .fap auxDecl.name #[] + return auxConst + +@[inline] +def castVarIfNeeded (fvarId : FVarId) (expectedType : Expr) (k : FVarId → BoxM (Code .impure)) : + BoxM (Code .impure) := do + let fvarIdType ← getType fvarId + if typesEqvForBoxing fvarIdType expectedType then + k fvarId + else + let v ← mkCast fvarId fvarIdType expectedType + let castDecl ← mkLetDecl .anonymous expectedType v + return .let castDecl (← k castDecl.fvarId) + +@[inline] +def castArgIfNeeded (arg : Arg .impure) (expectedType : Expr) + (k : Arg .impure → BoxM (Code .impure)) : BoxM (Code .impure) := do + match arg with + | .fvar fvarId => castVarIfNeeded fvarId expectedType (fun x => k (arg.updateFVar! x)) + | .erased => k arg + +def castArgsIfNeededAux (args : Array (Arg .impure)) (typeFromIdx : Nat → Expr) : + BoxM (Array (Arg .impure) × Array (CodeDecl .impure)) := do + let mut newArgs := Array.emptyWithCapacity args.size + let mut casters := #[] + for h : i in 0...args.size do + let arg := args[i] + let expectedType := typeFromIdx i + match arg with + | .erased => newArgs := newArgs.push arg + | .fvar fvarId => + let fvarIdType ← getType fvarId + if typesEqvForBoxing fvarIdType expectedType then + newArgs := newArgs.push arg + else + let v ← mkCast fvarId fvarIdType expectedType + let decl ← mkLetDecl .anonymous expectedType v + newArgs := newArgs.push <| .fvar decl.fvarId + casters := casters.push (.let decl) + return (newArgs, casters) + +@[inline] +def castArgsIfNeeded (args : Array (Arg .impure)) (ps : Array (Param .impure)) + (k : Array (Arg .impure) → BoxM (Code .impure)) : BoxM (Code .impure) := do + let (args, decls) ← castArgsIfNeededAux args fun i => ps[i]!.type + let k ← k args + return attachCodeDecls decls k + +@[inline] +def boxArgsIfNeeded (args : Array (Arg .impure)) (k : Array (Arg .impure) → BoxM (Code .impure)) : + BoxM (Code .impure) := do + let (args, decls) ← castArgsIfNeededAux args (fun _ => tobject) + let k ← k args + return attachCodeDecls decls k + +def unboxResultIfNeeded (code : Code .impure) (decl : LetDecl .impure) (k : Code .impure) : + BoxM (Code .impure) := do + if decl.type.isScalar then + let auxDecl ← mkLetDecl .anonymous tobject decl.value + let decl ← decl.updateValue (.unbox auxDecl.fvarId) + return .let auxDecl <| .let decl k + else + return code.updateLet! decl k + +def castResultIfNeeded (code : Code .impure) (decl : LetDecl .impure) (expType : Expr) + (k : Code .impure) : BoxM (Code .impure) := do + if typesEqvForBoxing decl.type expType then + return code.updateLet! decl k + else + let boxedTy := decl.type.boxed + let castDecl ← mkLetDecl .anonymous boxedTy decl.value + let castedValue ← mkCast castDecl.fvarId castDecl.type decl.type + let decl ← decl.updateValue castedValue + return .let castDecl <| .let decl k + + +/-- +Traverse `code`, trying to correct types through inference and ensuring that the ABI of other +functions is respected by inserting `box`/`unbox` operations. +-/ +partial def Code.explicitBoxing (code : Code .impure) : BoxM (Code .impure) := do + match code with + | .let decl k => visitLet code decl k + | .jp decl k => + let value ← decl.value.explicitBoxing + let decl ← decl.update (← getResultType) decl.params value + let k ← k.explicitBoxing + return code.updateFun! decl k + | .uset var i y k _ => + let k ← k.explicitBoxing + castVarIfNeeded y ImpureType.usize fun y => + return code.updateUset! var i y k + | .sset var i offset y ty k _ => + let k ← k.explicitBoxing + castVarIfNeeded y ty fun y => + return code.updateSset! var i offset y ty k + | .cases cs => + let alts ← cs.alts.mapMonoM (·.mapCodeM (·.explicitBoxing)) + castVarIfNeeded cs.discr (mkConst cs.typeName) fun discr => + return code.updateCases! (← getResultType) discr alts + | .return fvarId => + castVarIfNeeded fvarId (← getResultType) (fun fvarId => return code.updateReturn! fvarId) + | .jmp fvarId args => + let some jpDecl ← findFunDecl? fvarId | unreachable! + castArgsIfNeeded args jpDecl.params fun args => return code.updateJmp! fvarId args + | .unreach .. => return code.updateUnreach! (← getResultType) +where + /-- + Up to this point the type system of IR is quite loose so we can for example encounter situations + such as + ``` + let y : obj := f x + ``` + where `f : obj -> uint8`. It is the job of the boxing pass to enforce a stricter obj/scalar + separation and as such it needs to correct situations like this. + -/ + tryCorrectLetDeclType (currentType : Expr) (value : LetValue .impure) : BoxM Expr := do + match value with + | .fap f .. => + let some sig ← getImpureSignature? f | unreachable! + return sig.type + | .pap .. => return object + | .uproj .. => return usize + | .erased => return tagged + | .fvar .. | .lit .. | .sproj .. | .oproj .. | .reset .. | .ctor .. | .reuse .. => + return currentType + | .box .. | .unbox .. => unreachable! + + visitLet (code : Code .impure) (decl : LetDecl .impure) (k : Code .impure) : BoxM (Code .impure) := do + let type ← tryCorrectLetDeclType decl.type decl.value + let decl ← decl.update type decl.value + let k ← k.explicitBoxing + match decl.value with + | .ctor i args => + if i.isScalar && type.isScalar then + let decl ← decl.updateValue (.lit (.impureTypeScalarNumLit type i.cidx)) + return code.updateLet! decl k + else + boxArgsIfNeeded args fun args => do + let decl ← decl.updateValue (decl.value.updateArgs! args) + return code.updateLet! decl k + | .reuse _ _ _ args _ => + boxArgsIfNeeded args fun args => do + let decl ← decl.updateValue (decl.value.updateArgs! args) + return code.updateLet! decl k + | .fap f args => + let some sig ← getImpureSignature? f | unreachable! + castArgsIfNeeded args sig.params fun args => do + let decl ← decl.updateValue (decl.value.updateArgs! args) + castResultIfNeeded code decl sig.type k + | .pap f args => + let some sig ← getImpureSignature? f | unreachable! + let f := if ← requiresBoxedVersion sig then mkBoxedName f else f + boxArgsIfNeeded args fun args => do + let decl ← decl.updateValue (decl.value.updatePap! f args) + return code.updateLet! decl k + | .fvar _ args => + boxArgsIfNeeded args fun args => do + let decl ← decl.updateValue (decl.value.updateArgs! args) + unboxResultIfNeeded code decl k + | .erased | .reset .. | .sproj .. | .uproj .. | .oproj .. | .lit .. => + let decl ← decl.update type decl.value + return code.updateLet! decl k + | .box .. | .unbox .. => unreachable! + +def run (decls : Array (Decl .impure)) : CompilerM (Array (Decl .impure)) := do + let decls ← decls.foldlM (init := #[]) fun newDecls decl => do + match decl.value with + | .code code => + let s := { currDecl := decl.name, currDeclResultType := decl.type, decls } + let (code, s) ← code.explicitBoxing |>.run s |>.run {} + let newDecls := newDecls ++ s.auxDecls + let newDecl := { decl with value := .code code } + let newDecl ← newDecl.elimDeadVars + return newDecls.push newDecl + | .extern .. => return newDecls.push decl + addBoxedVersions decls + + +public def explicitBoxing : Pass where + phase := .impure + phaseOut := .impure + name := `boxing + run := run + +builtin_initialize + registerTraceClass `Compiler.explicitBoxing (inherited := true) + +end Lean.Compiler.LCNF diff --git a/src/Lean/Compiler/LCNF/FVarUtil.lean b/src/Lean/Compiler/LCNF/FVarUtil.lean index 8787dca138..7f0c1e67c6 100644 --- a/src/Lean/Compiler/LCNF/FVarUtil.lean +++ b/src/Lean/Compiler/LCNF/FVarUtil.lean @@ -81,12 +81,14 @@ def LetValue.mapFVarM [MonadLiftT CompilerM m] [Monad m] (f : FVarId → m FVarI | .reset n fvarId _ => return e.updateReset! n (← f fvarId) | .reuse fvarId i updateHeader args _ => return e.updateReuse! (← f fvarId) i updateHeader (← args.mapM (TraverseFVar.mapFVarM f)) + | .box ty fvarId _ => return e.updateBox! ty (← f fvarId) + | .unbox fvarId _ => return e.updateUnbox! (← f fvarId) def LetValue.forFVarM [Monad m] (f : FVarId → m Unit) (e : LetValue pu) : m Unit := do match e with | .lit .. | .erased => return () | .proj _ _ fvarId _ | .oproj _ fvarId _ | .sproj _ _ fvarId _ | .uproj _ fvarId _ - | .reset _ fvarId _ => f fvarId + | .reset _ fvarId _ | .box _ fvarId _ | .unbox fvarId _ => f fvarId | .const _ _ args _ | .pap _ args _ | .fap _ args _ | .ctor _ args _ => args.forM (TraverseFVar.forFVarM f) | .fvar fvarId args | .reuse fvarId _ _ args _ => f fvarId; args.forM (TraverseFVar.forFVarM f) diff --git a/src/Lean/Compiler/LCNF/Internalize.lean b/src/Lean/Compiler/LCNF/Internalize.lean index f0d6289c50..67224bf04b 100644 --- a/src/Lean/Compiler/LCNF/Internalize.lean +++ b/src/Lean/Compiler/LCNF/Internalize.lean @@ -112,7 +112,14 @@ private partial def internalizeLetValue (e : LetValue pu) : InternalizeM pu (Let match (← normFVar fvarId) with | .fvar fvarId' => return e.updateReuse! fvarId' info updateHeader (← internalizeArgs args) | .erased => return .erased - + | .unbox fvarId _ => + match (← normFVar fvarId) with + | .fvar fvarId' => return e.updateUnbox! fvarId' + | .erased => return .erased + | .box ty fvarId _ => + match (← normFVar fvarId) with + | .fvar fvarId' => return e.updateBox! ty fvarId' + | .erased => return .erased def internalizeLetDecl (decl : LetDecl pu) : InternalizeM pu (LetDecl pu) := do let binderName ← refreshBinderName decl.binderName diff --git a/src/Lean/Compiler/LCNF/Passes.lean b/src/Lean/Compiler/LCNF/Passes.lean index 652e08a500..a9182aff08 100644 --- a/src/Lean/Compiler/LCNF/Passes.lean +++ b/src/Lean/Compiler/LCNF/Passes.lean @@ -24,6 +24,7 @@ public import Lean.Compiler.LCNF.PushProj public import Lean.Compiler.LCNF.ResetReuse public import Lean.Compiler.LCNF.SimpCase public import Lean.Compiler.LCNF.InferBorrow +public import Lean.Compiler.LCNF.ExplicitBoxing public section @@ -149,6 +150,7 @@ def builtinPassManager : PassManager := { elimDeadVars (phase := .impure) (occurrence := 0), simpCase, inferBorrow, + explicitBoxing, inferVisibility (phase := .impure), saveImpure, -- End of impure phase ] diff --git a/src/Lean/Compiler/LCNF/PrettyPrinter.lean b/src/Lean/Compiler/LCNF/PrettyPrinter.lean index d1c5434949..49fd92bb76 100644 --- a/src/Lean/Compiler/LCNF/PrettyPrinter.lean +++ b/src/Lean/Compiler/LCNF/PrettyPrinter.lean @@ -92,6 +92,8 @@ def ppLetValue (e : LetValue pu) : M Format := do | .reset n fvarId _ => return f!"reset[{n}] {← ppFVar fvarId}" | .reuse fvarId info updateHeader args _ => return f!"reuse" ++ (if updateHeader then f!"!" else f!"") ++ f!" {← ppFVar fvarId} in {info}{← ppArgs args}" + | .box _ fvarId _ => return f!"box {← ppFVar fvarId}" + | .unbox fvarId _ => return f!"unbox {← ppFVar fvarId}" def ppParam (param : Param pu) : M Format := do let borrow := if param.borrow then "@&" else "" diff --git a/tests/lean/externBoxing.lean.expected.out b/tests/lean/externBoxing.lean.expected.out index 87ded6e43c..53a9c93aba 100644 --- a/tests/lean/externBoxing.lean.expected.out +++ b/tests/lean/externBoxing.lean.expected.out @@ -1,5 +1,5 @@ [Compiler.IR] [result] - extern _private.lean.externBoxing.0.Foo.bar (x_0 : obj) (x_1 : u64) : u64 + extern _private.lean.externBoxing.0.Foo.bar (x_1 : obj) (x_2 : u64) : u64 def _private.lean.externBoxing.0.Foo.bar._boxed (x_1 : obj) (x_2 : tobj) : tobj := let x_3 : u64 := unbox x_2; dec x_2; diff --git a/tests/lean/run/mangling.lean b/tests/lean/run/mangling.lean index 3ca056dd49..108c60154e 100644 --- a/tests/lean/run/mangling.lean +++ b/tests/lean/run/mangling.lean @@ -1,12 +1,15 @@ module import Lean.Compiler.IR.CompilerM import Lean.Compiler.NameMangling +import Lean.Compiler.LCNF.ExplicitBoxing /-! # Test behavior of name mangling -/ -open Lean IR ExplicitBoxing +open Lean IR +open Lean.Compiler.LCNF (mkBoxedName) + def checkMangle (n : Name) (s : String) : IO Unit := do if n.mangle "" ≠ s then throw <| .userError s!"failed: {n} mangles to {n.mangle ""} but expected {s}"