Documentation

Lean.Compiler.IR.Borrow

@[inline]

abbrev Lean.IR.Borrow.OwnedSet.Key :

Type

Equations

Lean.IR.Borrow.OwnedSet.Key = (Lean.IR.FunId × Lean.IR.Index)

def Lean.IR.Borrow.OwnedSet.beq :

Lean.IR.Borrow.OwnedSet.Key → Lean.IR.Borrow.OwnedSet.Key → Bool

Equations

Lean.IR.Borrow.OwnedSet.beq x x = match x, x with | (f₁, x₁), (f₂, x₂) => f₁ == f₂ && x₁ == x₂

instance Lean.IR.Borrow.OwnedSet.instBEqKey :

BEq Lean.IR.Borrow.OwnedSet.Key

Equations

Lean.IR.Borrow.OwnedSet.instBEqKey = { beq := Lean.IR.Borrow.OwnedSet.beq }

def Lean.IR.Borrow.OwnedSet.getHash :

Lean.IR.Borrow.OwnedSet.Key → UInt64

Equations

Lean.IR.Borrow.OwnedSet.getHash x = match x with | (f, x) => mixHash (hash f) (hash x)

instance Lean.IR.Borrow.OwnedSet.instHashableKey :

Hashable Lean.IR.Borrow.OwnedSet.Key

Equations

Lean.IR.Borrow.OwnedSet.instHashableKey = { hash := Lean.IR.Borrow.OwnedSet.getHash }

@[inline]

abbrev Lean.IR.Borrow.OwnedSet :

Type

Equations

Lean.IR.Borrow.OwnedSet = Std.HashMap Lean.IR.Borrow.OwnedSet.Key Unit

def Lean.IR.Borrow.OwnedSet.insert (s : Lean.IR.Borrow.OwnedSet) (k : Lean.IR.Borrow.OwnedSet.Key) :

Lean.IR.Borrow.OwnedSet

Equations

Lean.IR.Borrow.OwnedSet.insert s k = Std.HashMap.insert s k ()

def Lean.IR.Borrow.OwnedSet.contains (s : Lean.IR.Borrow.OwnedSet) (k : Lean.IR.Borrow.OwnedSet.Key) :

Equations

Lean.IR.Borrow.OwnedSet.contains s k = Std.HashMap.contains s k

inductive Lean.IR.Borrow.ParamMap.Key :

Type

decl: Lean.IR.FunId → Lean.IR.Borrow.ParamMap.Key
jp: Lean.IR.FunId → Lean.IR.JoinPointId → Lean.IR.Borrow.ParamMap.Key

instance Lean.IR.Borrow.ParamMap.instBEqKey :

BEq Lean.IR.Borrow.ParamMap.Key

Equations

Lean.IR.Borrow.ParamMap.instBEqKey = { beq := [anonymous] }

def Lean.IR.Borrow.ParamMap.getHash :

Lean.IR.Borrow.ParamMap.Key → UInt64

Equations

Lean.IR.Borrow.ParamMap.getHash x = match x with | Lean.IR.Borrow.ParamMap.Key.decl n => hash n | Lean.IR.Borrow.ParamMap.Key.jp n id => mixHash (hash n) (hash id)

instance Lean.IR.Borrow.ParamMap.instHashableKey :

Hashable Lean.IR.Borrow.ParamMap.Key

Equations

Lean.IR.Borrow.ParamMap.instHashableKey = { hash := Lean.IR.Borrow.ParamMap.getHash }

@[inline]

abbrev Lean.IR.Borrow.ParamMap :

Type

Equations

Lean.IR.Borrow.ParamMap = Std.HashMap Lean.IR.Borrow.ParamMap.Key (Array Lean.IR.Param)

def Lean.IR.Borrow.ParamMap.fmt (map : Lean.IR.Borrow.ParamMap) :

Equations

Lean.IR.Borrow.ParamMap.fmt map = let fmts := Std.HashMap.fold (fun fmt k ps => let k := match k with | Lean.IR.Borrow.ParamMap.Key.decl n => Lean.format n | Lean.IR.Borrow.ParamMap.Key.jp n id => Lean.format n ++ Std.Format.text ":" ++ Lean.format id; fmt ++ Lean.Format.line ++ k ++ Std.Format.text " -> " ++ Lean.IR.formatParams ps) Lean.Format.nil map; Std.Format.text "{" ++ Lean.Format.nest 1 fmts ++ Std.Format.text "}"

instance Lean.IR.Borrow.instToFormatParamMap :

Lean.ToFormat Lean.IR.Borrow.ParamMap

Equations

Lean.IR.Borrow.instToFormatParamMap = { format := Lean.IR.Borrow.ParamMap.fmt }

instance Lean.IR.Borrow.instToStringParamMap :

ToString Lean.IR.Borrow.ParamMap

Equations

Lean.IR.Borrow.instToStringParamMap = { toString := fun m => Lean.Format.pretty (Lean.format m) }

def Lean.IR.Borrow.InitParamMap.initBorrow (ps : Array Lean.IR.Param) :

Array Lean.IR.Param

Equations

Lean.IR.Borrow.InitParamMap.initBorrow ps = Array.map (fun p => { x := p.x, borrow := Lean.IR.IRType.isObj p.ty, ty := p.ty }) ps

def Lean.IR.Borrow.InitParamMap.initBorrowIfNotExported (exported : Bool) (ps : Array Lean.IR.Param) :

Array Lean.IR.Param

Equations

Lean.IR.Borrow.InitParamMap.initBorrowIfNotExported exported ps = if exported = true then ps else Lean.IR.Borrow.InitParamMap.initBorrow ps

partial def Lean.IR.Borrow.InitParamMap.visitFnBody (fnid : Lean.IR.FunId) :

Lean.IR.FnBody → StateM Lean.IR.Borrow.ParamMap Unit

def Lean.IR.Borrow.InitParamMap.visitDecls (env : Lean.Environment) (decls : Array Lean.IR.Decl) :

StateM Lean.IR.Borrow.ParamMap Unit

Equations

Lean.IR.Borrow.InitParamMap.visitDecls env decls = Array.forM (fun decl => match decl with | Lean.IR.Decl.fdecl f xs x b x_1 => let exported := Lean.isExport env f; do modify fun m => Std.HashMap.insert m (Lean.IR.Borrow.ParamMap.Key.decl f) (Lean.IR.Borrow.InitParamMap.initBorrowIfNotExported exported xs) Lean.IR.Borrow.InitParamMap.visitFnBody f b | x => pure ()) decls 0 (Array.size decls)

def Lean.IR.Borrow.mkInitParamMap (env : Lean.Environment) (decls : Array Lean.IR.Decl) :

Lean.IR.Borrow.ParamMap

Equations

Lean.IR.Borrow.mkInitParamMap env decls = StateT.run' (SeqRight.seqRight (Lean.IR.Borrow.InitParamMap.visitDecls env decls) fun x => get) ∅

partial def Lean.IR.Borrow.ApplyParamMap.visitFnBody (fn : Lean.IR.FunId) (paramMap : Lean.IR.Borrow.ParamMap) :

Lean.IR.FnBody → Lean.IR.FnBody

def Lean.IR.Borrow.ApplyParamMap.visitDecls (decls : Array Lean.IR.Decl) (paramMap : Lean.IR.Borrow.ParamMap) :

Array Lean.IR.Decl

Equations

Lean.IR.Borrow.ApplyParamMap.visitDecls decls paramMap = Array.map (fun decl => match decl with | Lean.IR.Decl.fdecl f xs ty b info => let b := Lean.IR.Borrow.ApplyParamMap.visitFnBody f paramMap b; match Std.HashMap.find? paramMap (Lean.IR.Borrow.ParamMap.Key.decl f) with | some xs => Lean.IR.Decl.fdecl f xs ty b info | none => panicWithPosWithDecl "Lean.Compiler.IR.Borrow" "Lean.IR.Borrow.ApplyParamMap.visitDecls" 129 19 "unreachable code has been reached" | other => other) decls

def Lean.IR.Borrow.applyParamMap (decls : Array Lean.IR.Decl) (map : Lean.IR.Borrow.ParamMap) :

Array Lean.IR.Decl

Equations

Lean.IR.Borrow.applyParamMap decls map = Lean.IR.Borrow.ApplyParamMap.visitDecls decls map

structure Lean.IR.Borrow.BorrowInfCtx :

Type

structure Lean.IR.Borrow.BorrowInfState :

Type

owned : Lean.IR.Borrow.OwnedSet
modified : Bool
paramMap : Lean.IR.Borrow.ParamMap

@[inline]

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

Type

Equations

Lean.IR.Borrow.M = ReaderT Lean.IR.Borrow.BorrowInfCtx (StateM Lean.IR.Borrow.BorrowInfState)

def Lean.IR.Borrow.getCurrFn :

Lean.IR.Borrow.M Lean.IR.FunId

Equations

Lean.IR.Borrow.getCurrFn = do let ctx ← read pure ctx.currFn

def Lean.IR.Borrow.markModified :

Lean.IR.Borrow.M Unit

Equations

Lean.IR.Borrow.markModified = modify fun s => { owned := s.owned, modified := true, paramMap := s.paramMap }

def Lean.IR.Borrow.ownVar (x : Lean.IR.VarId) :

Lean.IR.Borrow.M Unit

Equations

Lean.IR.Borrow.ownVar x = do let currFn ← Lean.IR.Borrow.getCurrFn modify fun s => if Lean.IR.Borrow.OwnedSet.contains s.owned (currFn, x.idx) = true then s else { owned := Lean.IR.Borrow.OwnedSet.insert s.owned (currFn, x.idx), modified := true, paramMap := s.paramMap }

def Lean.IR.Borrow.ownArg (x : Lean.IR.Arg) :

Lean.IR.Borrow.M Unit

Equations

Lean.IR.Borrow.ownArg x = match x with | Lean.IR.Arg.var x => Lean.IR.Borrow.ownVar x | x => pure ()

def Lean.IR.Borrow.ownArgs (xs : Array Lean.IR.Arg) :

Lean.IR.Borrow.M Unit

Equations

Lean.IR.Borrow.ownArgs xs = Array.forM Lean.IR.Borrow.ownArg xs 0 (Array.size xs)

def Lean.IR.Borrow.isOwned (x : Lean.IR.VarId) :

Lean.IR.Borrow.M Bool

Equations

Lean.IR.Borrow.isOwned x = do let currFn ← Lean.IR.Borrow.getCurrFn let s ← get pure (Lean.IR.Borrow.OwnedSet.contains s.owned (currFn, x.idx))

def Lean.IR.Borrow.updateParamMap (k : Lean.IR.Borrow.ParamMap.Key) :

Lean.IR.Borrow.M Unit

Equations

Lean.IR.Borrow.updateParamMap k = do let _ ← Lean.IR.Borrow.getCurrFn let s ← get match Std.HashMap.find? s.paramMap k with | some ps => do let ps ← Array.mapM (fun p => if (!p.borrow) = true then pure p else do let a ← Lean.IR.Borrow.isOwned p.x if a = true then do Lean.IR.Borrow.markModified pure { x := p.x, borrow := false, ty := p.ty } else pure p) ps modify fun s => { owned := s.owned, modified := s.modified, paramMap := Std.HashMap.insert s.paramMap k ps } | none => pure ()

def Lean.IR.Borrow.getParamInfo (k : Lean.IR.Borrow.ParamMap.Key) :

Lean.IR.Borrow.M (Array Lean.IR.Param)

Equations

Lean.IR.Borrow.getParamInfo k = do let s ← get match Std.HashMap.find? s.paramMap k with | some ps => pure ps | none => match k with | Lean.IR.Borrow.ParamMap.Key.decl fn => do let ctx ← read match Lean.IR.findEnvDecl ctx.env fn with | some decl => pure (Lean.IR.Decl.params decl) | none => panicWithPosWithDecl "Lean.Compiler.IR.Borrow" "Lean.IR.Borrow.getParamInfo" 203 21 "unreachable code has been reached" | x => panicWithPosWithDecl "Lean.Compiler.IR.Borrow" "Lean.IR.Borrow.getParamInfo" 204 11 "unreachable code has been reached"

def Lean.IR.Borrow.ownArgsUsingParams (xs : Array Lean.IR.Arg) (ps : Array Lean.IR.Param) :

Lean.IR.Borrow.M Unit

Equations

Lean.IR.Borrow.ownArgsUsingParams xs ps = Nat.forM (Array.size xs) fun i => let x := Array.getOp xs i; let p := Array.getOp ps i; if p.borrow = true then pure PUnit.unit else Lean.IR.Borrow.ownArg x

def Lean.IR.Borrow.ownParamsUsingArgs (xs : Array Lean.IR.Arg) (ps : Array Lean.IR.Param) :

Lean.IR.Borrow.M Unit

Equations

Lean.IR.Borrow.ownParamsUsingArgs xs ps = Nat.forM (Array.size xs) fun i => let x := Array.getOp xs i; let p := Array.getOp ps i; match x with | Lean.IR.Arg.var x => do let a ← Lean.IR.Borrow.isOwned x if a = true then Lean.IR.Borrow.ownVar p.x else pure PUnit.unit | x => pure ()

def Lean.IR.Borrow.ownArgsIfParam (xs : Array Lean.IR.Arg) :

Lean.IR.Borrow.M Unit

Equations

Lean.IR.Borrow.ownArgsIfParam xs = do let ctx ← read Array.forM (fun x => match x with | Lean.IR.Arg.var x => if Std.RBTree.contains ctx.paramSet x.idx = true then Lean.IR.Borrow.ownVar x else pure PUnit.unit | x => pure ()) xs 0 (Array.size xs)

def Lean.IR.Borrow.collectExpr (z : Lean.IR.VarId) :

Lean.IR.Expr → Lean.IR.Borrow.M Unit

Equations

Lean.IR.Borrow.collectExpr z x = match x with | Lean.IR.Expr.reset x x_1 => SeqRight.seqRight (Lean.IR.Borrow.ownVar z) fun x => Lean.IR.Borrow.ownVar x_1 | Lean.IR.Expr.reuse x x_1 x_2 ys => SeqRight.seqRight (SeqRight.seqRight (Lean.IR.Borrow.ownVar z) fun x_3 => Lean.IR.Borrow.ownVar x) fun x => Lean.IR.Borrow.ownArgsIfParam ys | Lean.IR.Expr.ctor x xs => SeqRight.seqRight (Lean.IR.Borrow.ownVar z) fun x => Lean.IR.Borrow.ownArgsIfParam xs | Lean.IR.Expr.proj x x_1 => do let a ← Lean.IR.Borrow.isOwned x_1 let _do_jp : Unit → Lean.IR.Borrow.M Unit := fun y => do let a ← Lean.IR.Borrow.isOwned z if a = true then Lean.IR.Borrow.ownVar x_1 else pure PUnit.unit if a = true then do let y ← Lean.IR.Borrow.ownVar z _do_jp y else do let y ← pure PUnit.unit _do_jp y | Lean.IR.Expr.fap g xs => do let ps ← Lean.IR.Borrow.getParamInfo (Lean.IR.Borrow.ParamMap.Key.decl g) SeqRight.seqRight (Lean.IR.Borrow.ownVar z) fun x => Lean.IR.Borrow.ownArgsUsingParams xs ps | Lean.IR.Expr.ap x ys => SeqRight.seqRight (SeqRight.seqRight (Lean.IR.Borrow.ownVar z) fun x_1 => Lean.IR.Borrow.ownVar x) fun x => Lean.IR.Borrow.ownArgs ys | Lean.IR.Expr.pap x xs => SeqRight.seqRight (Lean.IR.Borrow.ownVar z) fun x => Lean.IR.Borrow.ownArgs xs | other => pure ()

def Lean.IR.Borrow.preserveTailCall (x : Lean.IR.VarId) (v : Lean.IR.Expr) (b : Lean.IR.FnBody) :

Lean.IR.Borrow.M Unit

Equations

Lean.IR.Borrow.preserveTailCall x v b = do let ctx ← read match v, b with | Lean.IR.Expr.fap g ys, Lean.IR.FnBody.ret (Lean.IR.Arg.var z) => if (Array.any ctx.decls (fun a => Lean.IR.Decl.name a == g) 0 (Array.size ctx.decls) && x == z) = true then do let ps ← Lean.IR.Borrow.getParamInfo (Lean.IR.Borrow.ParamMap.Key.decl g) Lean.IR.Borrow.ownParamsUsingArgs ys ps else pure PUnit.unit | x, x_1 => pure ()

def Lean.IR.Borrow.updateParamSet (ctx : Lean.IR.Borrow.BorrowInfCtx) (ps : Array Lean.IR.Param) :

Lean.IR.Borrow.BorrowInfCtx

Equations

Lean.IR.Borrow.updateParamSet ctx ps = { env := ctx.env, decls := ctx.decls, currFn := ctx.currFn, paramSet := Array.foldl (fun s p => Std.RBTree.insert s p.x.idx) ctx.paramSet ps 0 (Array.size ps) }

partial def Lean.IR.Borrow.collectFnBody :

Lean.IR.FnBody → Lean.IR.Borrow.M Unit

def Lean.IR.Borrow.collectDecl :

Lean.IR.Decl → Lean.IR.Borrow.M Unit

Equations

Lean.IR.Borrow.collectDecl x = match x with | Lean.IR.Decl.fdecl f ys x b x_1 => withReader (fun ctx => let ctx := Lean.IR.Borrow.updateParamSet ctx ys; { env := ctx.env, decls := ctx.decls, currFn := f, paramSet := ctx.paramSet }) do Lean.IR.Borrow.collectFnBody b Lean.IR.Borrow.updateParamMap (Lean.IR.Borrow.ParamMap.Key.decl f) | x => pure ()

@[inline]

partial def Lean.IR.Borrow.whileModifing (x : Lean.IR.Borrow.M Unit) :

Lean.IR.Borrow.M Unit

def Lean.IR.Borrow.collectDecls :

Lean.IR.Borrow.M Lean.IR.Borrow.ParamMap

Equations

Lean.IR.Borrow.collectDecls = do let a ← read Lean.IR.Borrow.whileModifing (Array.forM Lean.IR.Borrow.collectDecl a.decls 0 (Array.size a.decls)) let s ← get pure s.paramMap

def Lean.IR.Borrow.infer (env : Lean.Environment) (decls : Array Lean.IR.Decl) :

Lean.IR.Borrow.ParamMap

Equations

Lean.IR.Borrow.infer env decls = StateT.run' (Lean.IR.Borrow.collectDecls { env := env, decls := decls, currFn := default, paramSet := ∅ }) { owned := ∅, modified := false, paramMap := Lean.IR.Borrow.mkInitParamMap env decls }

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

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

Equations

Lean.IR.inferBorrow decls = do let env ← Lean.IR.getEnv let paramMap : Lean.IR.Borrow.ParamMap := Lean.IR.Borrow.infer env decls pure (Lean.IR.Borrow.applyParamMap decls paramMap)