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.

@[reducible, inline]

abbrev Lean.IR.ExplicitBoxing.N (α : Type) : Type

Equations

• Lean.IR.ExplicitBoxing.N = StateM Nat

def Lean.IR.ExplicitBoxing.requiresBoxedVersion (env : Environment) (decl : Decl) : Bool

Equations

• One or more equations did not get rendered due to their size.

def Lean.IR.ExplicitBoxing.mkBoxedVersionAux (decl : Decl) : N Decl

Equations

• One or more equations did not get rendered due to their size.

def Lean.IR.ExplicitBoxing.mkBoxedVersion (decl : Decl) : Decl

Equations

• Lean.IR.ExplicitBoxing.mkBoxedVersion decl = StateT.run' (Lean.IR.ExplicitBoxing.mkBoxedVersionAux decl) 1

def Lean.IR.ExplicitBoxing.addBoxedVersions (env : Environment) (decls : Array Decl) : Array Decl

Equations

• One or more equations did not get rendered due to their size.

def Lean.IR.ExplicitBoxing.eqvTypes (t₁ t₂ : IRType) : Bool

Equations

• Lean.IR.ExplicitBoxing.eqvTypes t₁ t₂ = (t₁.isScalar == t₂.isScalar && (!t₁.isScalar || t₁ == t₂))

structure Lean.IR.ExplicitBoxing.BoxingContext : Type

• f : FunId
• localCtx : LocalContext
• resultType : IRType
• decls : Array Decl
• env : Environment

structure Lean.IR.ExplicitBoxing.BoxingState : Type

• nextIdx : Index
• auxDecls : Array Decl

  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.

• auxDeclCache : AssocList FnBody Expr
• nextAuxId : Nat

@[reducible, inline]

abbrev Lean.IR.ExplicitBoxing.M (α : Type) : Type

Equations

• Lean.IR.ExplicitBoxing.M = ReaderT Lean.IR.ExplicitBoxing.BoxingContext (StateT Lean.IR.ExplicitBoxing.BoxingState Id)

def Lean.IR.ExplicitBoxing.getEnv : M Environment

Equations

• Lean.IR.ExplicitBoxing.getEnv = Lean.IR.ExplicitBoxing.BoxingContext.env <$> read

def Lean.IR.ExplicitBoxing.getLocalContext : M LocalContext

Equations

• Lean.IR.ExplicitBoxing.getLocalContext = Lean.IR.ExplicitBoxing.BoxingContext.localCtx <$> read

def Lean.IR.ExplicitBoxing.getResultType : M IRType

Equations

• Lean.IR.ExplicitBoxing.getResultType = Lean.IR.ExplicitBoxing.BoxingContext.resultType <$> read

def Lean.IR.ExplicitBoxing.getVarType (x : VarId) : M IRType

Equations

• Lean.IR.ExplicitBoxing.getVarType x = do let localCtx ← Lean.IR.ExplicitBoxing.getLocalContext match localCtx.getType x with | some t => pure t | none => pure Lean.IR.IRType.tobject

def Lean.IR.ExplicitBoxing.getJPParams (j : JoinPointId) : M (Array Param)

Equations

• Lean.IR.ExplicitBoxing.getJPParams j = do let localCtx ← Lean.IR.ExplicitBoxing.getLocalContext match localCtx.getJPParams j with | some ys => pure ys | none => pure #[]

def Lean.IR.ExplicitBoxing.getDecl (fid : FunId) : M Decl

Equations

• Lean.IR.ExplicitBoxing.getDecl fid = do let ctx ← read match Lean.IR.findEnvDecl' ctx.env fid ctx.decls with | some decl => pure decl | none => pure default

@[inline]

def Lean.IR.ExplicitBoxing.withParams {α : Type} (xs : Array Param) (k : M α) : M α

Equations

• One or more equations did not get rendered due to their size.

@[inline]

def Lean.IR.ExplicitBoxing.withVDecl {α : Type} (x : VarId) (ty : IRType) (v : Expr) (k : M α) : M α

Equations

• One or more equations did not get rendered due to their size.

@[inline]

def Lean.IR.ExplicitBoxing.withJDecl {α : Type} (j : JoinPointId) (xs : Array Param) (v : FnBody) (k : M α) : M α

Equations

• One or more equations did not get rendered due to their size.

def Lean.IR.ExplicitBoxing.mkCast (x : VarId) (xType expectedType : IRType) : M Expr

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.

Equations

• One or more equations did not get rendered due to their size.

@[inline]

def Lean.IR.ExplicitBoxing.castVarIfNeeded (x : VarId) (expected : IRType) (k : VarId → M FnBody) : M FnBody

Equations

• One or more equations did not get rendered due to their size.

@[inline]

def Lean.IR.ExplicitBoxing.castArgIfNeeded (x : Arg) (expected : IRType) (k : Arg → M FnBody) : M FnBody

Equations

• Lean.IR.ExplicitBoxing.castArgIfNeeded (Lean.IR.Arg.var x_2) expected k = Lean.IR.ExplicitBoxing.castVarIfNeeded x_2 expected fun (x : Lean.IR.VarId) => k (Lean.IR.Arg.var x)
• Lean.IR.ExplicitBoxing.castArgIfNeeded Lean.IR.Arg.erased expected k = k Lean.IR.Arg.erased

def Lean.IR.ExplicitBoxing.castArgsIfNeededAux (xs : Array Arg) (typeFromIdx : Nat → IRType) : M (Array Arg × Array FnBody)

Equations

• One or more equations did not get rendered due to their size.

@[inline]

def Lean.IR.ExplicitBoxing.castArgsIfNeeded (xs : Array Arg) (ps : Array Param) (k : Array Arg → M FnBody) : M FnBody

Equations

• One or more equations did not get rendered due to their size.

@[inline]

def Lean.IR.ExplicitBoxing.boxArgsIfNeeded (xs : Array Arg) (k : Array Arg → M FnBody) : M FnBody

Equations

• One or more equations did not get rendered due to their size.

def Lean.IR.ExplicitBoxing.unboxResultIfNeeded (x : VarId) (ty : IRType) (e : Expr) (b : FnBody) : M FnBody

Equations

• One or more equations did not get rendered due to their size.

def Lean.IR.ExplicitBoxing.castResultIfNeeded (x : VarId) (ty : IRType) (e : Expr) (eType : IRType) (b : FnBody) : M FnBody

Equations

• One or more equations did not get rendered due to their size.

def Lean.IR.ExplicitBoxing.visitVDeclExpr (x : VarId) (ty : IRType) (e : Expr) (b : FnBody) : M FnBody

Equations

• One or more equations did not get rendered due to their size.
• Lean.IR.ExplicitBoxing.visitVDeclExpr x ty e b = pure (Lean.IR.FnBody.vdecl x ty e b)

def Lean.IR.ExplicitBoxing.tryCorrectVDeclType (ty : IRType) (e : Expr) : M IRType

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.

Equations

• One or more equations did not get rendered due to their size.
• Lean.IR.ExplicitBoxing.tryCorrectVDeclType ty (Lean.IR.Expr.fap f ys) = do let decl ← Lean.IR.ExplicitBoxing.getDecl f pure decl.resultType
• Lean.IR.ExplicitBoxing.tryCorrectVDeclType ty (Lean.IR.Expr.pap c ys) = pure Lean.IR.IRType.object
• Lean.IR.ExplicitBoxing.tryCorrectVDeclType ty (Lean.IR.Expr.uproj i x) = pure Lean.IR.IRType.usize
• Lean.IR.ExplicitBoxing.tryCorrectVDeclType ty (Lean.IR.Expr.ctor i ys) = pure ty
• Lean.IR.ExplicitBoxing.tryCorrectVDeclType ty (Lean.IR.Expr.reuse x i updtHeader ys) = pure ty
• Lean.IR.ExplicitBoxing.tryCorrectVDeclType ty (Lean.IR.Expr.ap x ys) = pure ty
• Lean.IR.ExplicitBoxing.tryCorrectVDeclType ty (Lean.IR.Expr.lit v) = pure ty
• Lean.IR.ExplicitBoxing.tryCorrectVDeclType ty (Lean.IR.Expr.sproj n offset x) = pure ty
• Lean.IR.ExplicitBoxing.tryCorrectVDeclType ty (Lean.IR.Expr.proj i x) = pure ty
• Lean.IR.ExplicitBoxing.tryCorrectVDeclType ty (Lean.IR.Expr.reset n x) = pure ty

partial def Lean.IR.ExplicitBoxing.visitFnBody : FnBody → M FnBody

def Lean.IR.ExplicitBoxing.run (env : Environment) (decls : Array Decl) : Array Decl

Equations

• One or more equations did not get rendered due to their size.

def Lean.IR.explicitBoxing (decls : Array Decl) : CompilerM (Array Decl)

Equations

• Lean.IR.explicitBoxing decls = do let env ← Lean.getEnv pure (Lean.IR.ExplicitBoxing.run env decls)