Documentation

Lean.Compiler.IR.ExpandResetReuse

@[inline]

abbrev Lean.IR.ExpandResetReuse.ProjMap :

Type

Equations

Lean.IR.ExpandResetReuse.ProjMap = Std.HashMap Lean.IR.VarId Lean.IR.Expr

@[inline]

abbrev Lean.IR.ExpandResetReuse.CollectProjMap.Collector :

Type

Equations

Lean.IR.ExpandResetReuse.CollectProjMap.Collector = (Lean.IR.ExpandResetReuse.ProjMap → Lean.IR.ExpandResetReuse.ProjMap)

@[inline]

def Lean.IR.ExpandResetReuse.CollectProjMap.collectVDecl (x : Lean.IR.VarId) (v : Lean.IR.Expr) :

Lean.IR.ExpandResetReuse.CollectProjMap.Collector

Equations

Lean.IR.ExpandResetReuse.CollectProjMap.collectVDecl x v m = match v with | Lean.IR.Expr.proj x x_1 => Std.HashMap.insert m x v | Lean.IR.Expr.sproj x x_1 x_2 => Std.HashMap.insert m x v | Lean.IR.Expr.uproj x x_1 => Std.HashMap.insert m x v | x => m

partial def Lean.IR.ExpandResetReuse.CollectProjMap.collectFnBody :

Lean.IR.FnBody → Lean.IR.ExpandResetReuse.CollectProjMap.Collector

def Lean.IR.ExpandResetReuse.mkProjMap (d : Lean.IR.Decl) :

Lean.IR.ExpandResetReuse.ProjMap

Equations

Lean.IR.ExpandResetReuse.mkProjMap d = match d with | Lean.IR.Decl.fdecl x x_1 x_2 b x_3 => Lean.IR.ExpandResetReuse.CollectProjMap.collectFnBody b ∅ | x => ∅

structure Lean.IR.ExpandResetReuse.Context :

Type

projMap : Lean.IR.ExpandResetReuse.ProjMap

partial def Lean.IR.ExpandResetReuse.consumed (x : Lean.IR.VarId) :

Lean.IR.FnBody → Bool

@[inline]

abbrev Lean.IR.ExpandResetReuse.Mask :

Type

Equations

Lean.IR.ExpandResetReuse.Mask = Array (Option Lean.IR.VarId)

partial def Lean.IR.ExpandResetReuse.eraseProjIncForAux (y : Lean.IR.VarId) (bs : Array Lean.IR.FnBody) (mask : Lean.IR.ExpandResetReuse.Mask) (keep : Array Lean.IR.FnBody) :

Array Lean.IR.FnBody × Lean.IR.ExpandResetReuse.Mask

def Lean.IR.ExpandResetReuse.eraseProjIncFor (n : Nat) (y : Lean.IR.VarId) (bs : Array Lean.IR.FnBody) :

Array Lean.IR.FnBody × Lean.IR.ExpandResetReuse.Mask

Equations

Lean.IR.ExpandResetReuse.eraseProjIncFor n y bs = Lean.IR.ExpandResetReuse.eraseProjIncForAux y bs (mkArray n none) #[]

partial def Lean.IR.ExpandResetReuse.reuseToCtor (x : Lean.IR.VarId) :

Lean.IR.FnBody → Lean.IR.FnBody

def Lean.IR.ExpandResetReuse.mkSlowPath (x : Lean.IR.VarId) (y : Lean.IR.VarId) (mask : Lean.IR.ExpandResetReuse.Mask) (b : Lean.IR.FnBody) :

Equations

Lean.IR.ExpandResetReuse.mkSlowPath x y mask b = let b := Lean.IR.ExpandResetReuse.reuseToCtor x b; let b := Lean.IR.FnBody.dec y 1 true false b; Array.foldl (fun b m => match m with | some z => Lean.IR.FnBody.inc z 1 true false b | none => b) b mask 0 (Array.size mask)

@[inline]

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

Type

Equations

Lean.IR.ExpandResetReuse.M = ReaderT Lean.IR.ExpandResetReuse.Context (StateM Nat)

def Lean.IR.ExpandResetReuse.mkFresh :

Lean.IR.ExpandResetReuse.M Lean.IR.VarId

Equations

Lean.IR.ExpandResetReuse.mkFresh = modifyGet fun n => ({ idx := n }, n + 1)

def Lean.IR.ExpandResetReuse.releaseUnreadFields (y : Lean.IR.VarId) (mask : Lean.IR.ExpandResetReuse.Mask) (b : Lean.IR.FnBody) :

Lean.IR.ExpandResetReuse.M Lean.IR.FnBody

Equations

Lean.IR.ExpandResetReuse.releaseUnreadFields y mask b = Nat.foldM (fun i b => match Array.get! mask i with | some x => pure b | none => do let fld ← Lean.IR.ExpandResetReuse.mkFresh pure (Lean.IR.FnBody.vdecl fld Lean.IR.IRType.object (Lean.IR.Expr.proj i y) (Lean.IR.FnBody.dec fld 1 true false b))) b (Array.size mask)

def Lean.IR.ExpandResetReuse.setFields (y : Lean.IR.VarId) (zs : Array Lean.IR.Arg) (b : Lean.IR.FnBody) :

Equations

Lean.IR.ExpandResetReuse.setFields y zs b = Nat.fold (fun i b => Lean.IR.FnBody.set y i (Array.get! zs i) b) (Array.size zs) b

def Lean.IR.ExpandResetReuse.isSelfSet (ctx : Lean.IR.ExpandResetReuse.Context) (x : Lean.IR.VarId) (i : Nat) (y : Lean.IR.Arg) :

Equations

Lean.IR.ExpandResetReuse.isSelfSet ctx x i y = match y with | Lean.IR.Arg.var y => match Std.HashMap.find? ctx.projMap y with | some (Lean.IR.Expr.proj j w) => j == i && w == x | x => false | x => false

def Lean.IR.ExpandResetReuse.isSelfUSet (ctx : Lean.IR.ExpandResetReuse.Context) (x : Lean.IR.VarId) (i : Nat) (y : Lean.IR.VarId) :

Equations

Lean.IR.ExpandResetReuse.isSelfUSet ctx x i y = match Std.HashMap.find? ctx.projMap y with | some (Lean.IR.Expr.uproj j w) => j == i && w == x | x => false

def Lean.IR.ExpandResetReuse.isSelfSSet (ctx : Lean.IR.ExpandResetReuse.Context) (x : Lean.IR.VarId) (n : Nat) (i : Nat) (y : Lean.IR.VarId) :

Equations

Lean.IR.ExpandResetReuse.isSelfSSet ctx x n i y = match Std.HashMap.find? ctx.projMap y with | some (Lean.IR.Expr.sproj m j w) => n == m && j == i && w == x | x => false

partial def Lean.IR.ExpandResetReuse.removeSelfSet (ctx : Lean.IR.ExpandResetReuse.Context) :

Lean.IR.FnBody → Lean.IR.FnBody

partial def Lean.IR.ExpandResetReuse.reuseToSet (ctx : Lean.IR.ExpandResetReuse.Context) (x : Lean.IR.VarId) (y : Lean.IR.VarId) :

Lean.IR.FnBody → Lean.IR.FnBody

def Lean.IR.ExpandResetReuse.mkFastPath (x : Lean.IR.VarId) (y : Lean.IR.VarId) (mask : Lean.IR.ExpandResetReuse.Mask) (b : Lean.IR.FnBody) :

Lean.IR.ExpandResetReuse.M Lean.IR.FnBody

Equations

Lean.IR.ExpandResetReuse.mkFastPath x y mask b = do let ctx ← read let b : Lean.IR.FnBody := Lean.IR.ExpandResetReuse.reuseToSet ctx x y b Lean.IR.ExpandResetReuse.releaseUnreadFields y mask b

def Lean.IR.ExpandResetReuse.expand (mainFn : Lean.IR.FnBody → Array Lean.IR.FnBody → Lean.IR.ExpandResetReuse.M Lean.IR.FnBody) (bs : Array Lean.IR.FnBody) (x : Lean.IR.VarId) (n : Nat) (y : Lean.IR.VarId) (b : Lean.IR.FnBody) :

Lean.IR.ExpandResetReuse.M Lean.IR.FnBody

Equations

Lean.IR.ExpandResetReuse.expand mainFn bs x n y b = let bOld := Lean.IR.FnBody.vdecl x Lean.IR.IRType.object (Lean.IR.Expr.reset n y) b; let x := Lean.IR.ExpandResetReuse.eraseProjIncFor n y bs; match x with | (bs, mask) => let bSlow := Lean.IR.ExpandResetReuse.mkSlowPath x y mask b; do let bFast ← Lean.IR.ExpandResetReuse.mkFastPath x y mask b let bFast ← mainFn bFast #[] let c ← Lean.IR.ExpandResetReuse.mkFresh let b : Lean.IR.FnBody := Lean.IR.FnBody.vdecl c Lean.IR.IRType.uint8 (Lean.IR.Expr.isShared y) (Lean.IR.mkIf c bSlow bFast) pure (Lean.IR.reshape bs b)

partial def Lean.IR.ExpandResetReuse.searchAndExpand :

Lean.IR.FnBody → Array Lean.IR.FnBody → Lean.IR.ExpandResetReuse.M Lean.IR.FnBody

def Lean.IR.ExpandResetReuse.main (d : Lean.IR.Decl) :

Equations

Lean.IR.ExpandResetReuse.main d = match d with | Lean.IR.Decl.fdecl x x_1 x_2 b x_3 => let m := Lean.IR.ExpandResetReuse.mkProjMap d; let nextIdx := Lean.IR.Decl.maxIndex d + 1; let bNew := StateT.run' (Lean.IR.ExpandResetReuse.searchAndExpand b #[] { projMap := m }) nextIdx; Lean.IR.Decl.updateBody! d bNew | d => d

def Lean.IR.Decl.expandResetReuse (d : Lean.IR.Decl) :

Equations

Lean.IR.Decl.expandResetReuse d = Lean.IR.Decl.normalizeIds (Lean.IR.ExpandResetReuse.main d)