Lean.Compiler.IR.Boxing

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

source

@[inline]

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

Lean.IR.ExplicitBoxing.M α

Equations

Lean.IR.ExplicitBoxing.withParams xs k = withReader (fun ctx => { f := ctx.f, localCtx := Lean.IR.LocalContext.addParams ctx.localCtx xs, resultType := ctx.resultType, decls := ctx.decls, env := ctx.env }) k

source

@[inline]

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

Lean.IR.ExplicitBoxing.M α

Equations

Lean.IR.ExplicitBoxing.withVDecl x ty v k = withReader (fun ctx => { f := ctx.f, localCtx := Lean.IR.LocalContext.addLocal ctx.localCtx x ty v, resultType := ctx.resultType, decls := ctx.decls, env := ctx.env }) k

source

@[inline]

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

Lean.IR.ExplicitBoxing.M α

Equations

Lean.IR.ExplicitBoxing.withJDecl j xs v k = withReader (fun ctx => { f := ctx.f, localCtx := Lean.IR.LocalContext.addJP ctx.localCtx j xs v, resultType := ctx.resultType, decls := ctx.decls, env := ctx.env }) k

source

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

Lean.IR.ExplicitBoxing.M Lean.IR.Expr

Equations

Lean.IR.ExplicitBoxing.mkCast x xType expectedType = do let a ← Lean.IR.ExplicitBoxing.isExpensiveConstantValueBoxing x xType match a with | some v => do let ctx ← read let s ← get let body : Lean.IR.FnBody := Lean.IR.FnBody.vdecl { idx := 1 } xType v (Lean.IR.FnBody.vdecl { idx := 2 } expectedType (Lean.IR.Expr.box xType { idx := 1 }) (Lean.IR.FnBody.ret (Lean.IR.mkVarArg { idx := 2 }))) match Std.AssocList.find? body s.auxDeclCache with | some v => pure v | none => let auxName := ctx.f ++ Lean.Name.appendIndexAfter `_boxed_const s.nextAuxId; let auxConst := Lean.IR.Expr.fap auxName #[]; let auxDecl := Lean.IR.Decl.fdecl auxName #[] expectedType body { sorryDep? := none }; do modify fun s => { nextIdx := s.nextIdx, auxDecls := Array.push s.auxDecls auxDecl, auxDeclCache := Std.AssocList.cons body auxConst s.auxDeclCache, nextAuxId := s.nextAuxId + 1 } pure auxConst | none => pure (if Lean.IR.IRType.isScalar xType = true then Lean.IR.Expr.box xType x else Lean.IR.Expr.unbox x)

source

@[inline]

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

Lean.IR.ExplicitBoxing.M Lean.IR.FnBody

Equations

Lean.IR.ExplicitBoxing.castVarIfNeeded x expected k = do let xType ← Lean.IR.ExplicitBoxing.getVarType x if Lean.IR.ExplicitBoxing.eqvTypes xType expected = true then k x else do let y ← Lean.IR.ExplicitBoxing.M.mkFresh let v ← Lean.IR.ExplicitBoxing.mkCast x xType expected Lean.IR.FnBody.vdecl y expected v <$> k y

source

@[inline]

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

Lean.IR.ExplicitBoxing.M Lean.IR.FnBody

Equations

Lean.IR.ExplicitBoxing.castArgIfNeeded x expected k = match x with | Lean.IR.Arg.var x => Lean.IR.ExplicitBoxing.castVarIfNeeded x expected fun x => k (Lean.IR.Arg.var x) | x => k x

source

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

Lean.IR.ExplicitBoxing.M (Array Lean.IR.Arg × Array Lean.IR.FnBody)

Equations

Lean.IR.ExplicitBoxing.castArgsIfNeededAux xs typeFromIdx = let xs' := #[]; let bs := #[]; let i := 0; do let r ← forIn xs { fst := xs', snd := { fst := bs, snd := i } } fun x r => let xs' := r.fst; let x_1 := r.snd; let bs := x_1.fst; let i := x_1.snd; let expected := typeFromIdx i; let _do_jp := fun xs' bs i y => let i := i + 1; do pure PUnit.unit pure (ForInStep.yield { fst := xs', snd := { fst := bs, snd := i } }); match x with | Lean.IR.Arg.irrelevant => let xs' := Array.push xs' x; do let y ← pure PUnit.unit _do_jp xs' bs i y | Lean.IR.Arg.var x => do let xType ← Lean.IR.ExplicitBoxing.getVarType x if Lean.IR.ExplicitBoxing.eqvTypes xType expected = true then let xs' := Array.push xs' (Lean.IR.Arg.var x); do let y ← pure PUnit.unit _do_jp xs' bs i y else do let y ← Lean.IR.ExplicitBoxing.M.mkFresh let v ← Lean.IR.ExplicitBoxing.mkCast x xType expected let b : Lean.IR.FnBody := Lean.IR.FnBody.vdecl y expected v Lean.IR.FnBody.nil let xs' : Array Lean.IR.Arg := Array.push xs' (Lean.IR.Arg.var y) let bs : Array Lean.IR.FnBody := Array.push bs b let y ← pure PUnit.unit _do_jp xs' bs i y let x : MProd (Array Lean.IR.Arg) (MProd (Array Lean.IR.FnBody) Nat) := r match x with | { fst := xs', snd := { fst := bs, snd := i } } => pure (xs', bs)

source

@[inline]

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

Lean.IR.ExplicitBoxing.M Lean.IR.FnBody

Equations

Lean.IR.ExplicitBoxing.castArgsIfNeeded xs ps k = do let discr ← Lean.IR.ExplicitBoxing.castArgsIfNeededAux xs fun i => (Array.getOp ps i).ty let x : Array Lean.IR.Arg × Array Lean.IR.FnBody := discr match x with | (ys, bs) => do let b ← k ys pure (Lean.IR.reshape bs b)

source

@[inline]

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

Lean.IR.ExplicitBoxing.M Lean.IR.FnBody

Equations

Lean.IR.ExplicitBoxing.boxArgsIfNeeded xs k = do let discr ← Lean.IR.ExplicitBoxing.castArgsIfNeededAux xs fun x => Lean.IR.IRType.object let x : Array Lean.IR.Arg × Array Lean.IR.FnBody := discr match x with | (ys, bs) => do let b ← k ys pure (Lean.IR.reshape bs b)

source

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

Lean.IR.ExplicitBoxing.M Lean.IR.FnBody

Equations

Lean.IR.ExplicitBoxing.unboxResultIfNeeded x ty e b = if Lean.IR.IRType.isScalar ty = true then do let y ← Lean.IR.ExplicitBoxing.M.mkFresh pure (Lean.IR.FnBody.vdecl y Lean.IR.IRType.object e (Lean.IR.FnBody.vdecl x ty (Lean.IR.Expr.unbox y) b)) else pure (Lean.IR.FnBody.vdecl x ty e b)

source

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

Lean.IR.ExplicitBoxing.M Lean.IR.FnBody

Equations

Lean.IR.ExplicitBoxing.castResultIfNeeded x ty e eType b = if Lean.IR.ExplicitBoxing.eqvTypes ty eType = true then pure (Lean.IR.FnBody.vdecl x ty e b) else do let y ← Lean.IR.ExplicitBoxing.M.mkFresh let v ← Lean.IR.ExplicitBoxing.mkCast y eType ty pure (Lean.IR.FnBody.vdecl y eType e (Lean.IR.FnBody.vdecl x ty v b))

source

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

Lean.IR.ExplicitBoxing.M Lean.IR.FnBody

Equations

Lean.IR.ExplicitBoxing.visitVDeclExpr x ty e b = match e with | Lean.IR.Expr.ctor c ys => if (Lean.IR.CtorInfo.isScalar c && Lean.IR.IRType.isScalar ty) = true then pure (Lean.IR.FnBody.vdecl x ty (Lean.IR.Expr.lit (Lean.IR.LitVal.num c.cidx)) b) else Lean.IR.ExplicitBoxing.boxArgsIfNeeded ys fun ys => pure (Lean.IR.FnBody.vdecl x ty (Lean.IR.Expr.ctor c ys) b) | Lean.IR.Expr.reuse w c u ys => Lean.IR.ExplicitBoxing.boxArgsIfNeeded ys fun ys => pure (Lean.IR.FnBody.vdecl x ty (Lean.IR.Expr.reuse w c u ys) b) | Lean.IR.Expr.fap f ys => do let decl ← Lean.IR.ExplicitBoxing.getDecl f Lean.IR.ExplicitBoxing.castArgsIfNeeded ys (Lean.IR.Decl.params decl) fun ys => Lean.IR.ExplicitBoxing.castResultIfNeeded x ty (Lean.IR.Expr.fap f ys) (Lean.IR.Decl.resultType decl) b | Lean.IR.Expr.pap f ys => do let env ← Lean.IR.ExplicitBoxing.getEnv let decl ← Lean.IR.ExplicitBoxing.getDecl f let f : Lean.Name := if Lean.IR.ExplicitBoxing.requiresBoxedVersion env decl = true then Lean.IR.ExplicitBoxing.mkBoxedName f else f Lean.IR.ExplicitBoxing.boxArgsIfNeeded ys fun ys => pure (Lean.IR.FnBody.vdecl x ty (Lean.IR.Expr.pap f ys) b) | Lean.IR.Expr.ap f ys => Lean.IR.ExplicitBoxing.boxArgsIfNeeded ys fun ys => Lean.IR.ExplicitBoxing.unboxResultIfNeeded x ty (Lean.IR.Expr.ap f ys) b | other => pure (Lean.IR.FnBody.vdecl x ty e b)

source

partial def Lean.IR.ExplicitBoxing.visitFnBody :

Lean.IR.FnBody → Lean.IR.ExplicitBoxing.M Lean.IR.FnBody

source

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

Array Lean.IR.Decl

Equations

Lean.IR.ExplicitBoxing.run env decls = let ctx := { f := default, localCtx := ∅, resultType := Lean.IR.IRType.irrelevant, decls := decls, env := env }; let decls := Array.foldl (fun newDecls decl => match decl with | Lean.IR.Decl.fdecl f xs t b x => let nextIdx := Lean.IR.Decl.maxIndex decl + 1; let x := StateT.run (Lean.IR.ExplicitBoxing.withParams xs (Lean.IR.ExplicitBoxing.visitFnBody b) { f := f, localCtx := ctx.localCtx, resultType := t, decls := ctx.decls, env := ctx.env }) { nextIdx := nextIdx, auxDecls := #[], auxDeclCache := Std.AssocList.empty, nextAuxId := 1 }; match x with | (b, s) => let newDecls := newDecls ++ s.auxDecls; let newDecl := Lean.IR.Decl.updateBody! decl b; let newDecl := Lean.IR.Decl.elimDead newDecl; Array.push newDecls newDecl | d => Array.push newDecls d) #[] decls 0 (Array.size decls); Lean.IR.ExplicitBoxing.addBoxedVersions env decls

source

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

Lean.IR.CompilerM (Array Lean.IR.Decl)

Equations

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