151 lines
5.5 KiB
Text
151 lines
5.5 KiB
Text
/-
|
||
Copyright (c) 2022 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.Compiler.LCNF.Basic
|
||
import Init.Omega
|
||
|
||
public section
|
||
|
||
namespace Lean.Compiler.LCNF
|
||
|
||
namespace ToExpr
|
||
|
||
abbrev LevelMap := FVarIdMap Nat
|
||
|
||
private def _root_.Lean.FVarId.toExpr (offset : Nat) (m : LevelMap) (fvarId : FVarId) : Expr :=
|
||
match m.get? fvarId with
|
||
| some level => .bvar (offset - level - 1)
|
||
| none => .fvar fvarId
|
||
|
||
private def _root_.Lean.Expr.abstract' (offset : Nat) (m : LevelMap) (e : Expr) : Expr :=
|
||
go offset e
|
||
where
|
||
go (o : Nat) (e : Expr) : Expr :=
|
||
match e with
|
||
| .fvar fvarId => fvarId.toExpr o m
|
||
| .lit .. | .const .. | .sort .. | .mvar .. | .bvar .. => e
|
||
| .app f a => .app (go o f) (go o a)
|
||
| .mdata k b => .mdata k (go o b)
|
||
| .proj s i b => .proj s i (go o b)
|
||
| .forallE n d b bi => .forallE n (go o d) (go (o+1) b) bi
|
||
| .lam n d b bi => .lam n (go o d) (go (o+1) b) bi
|
||
| .letE n t v b nd => .letE n (go o t) (go o v) (go (o+1) b) nd
|
||
|
||
abbrev ToExprM := ReaderT Nat $ StateM LevelMap
|
||
|
||
@[inline] def mkLambdaM (params : Array (Param pu)) (e : Expr) : ToExprM Expr :=
|
||
return go (← read) (← get) params.size e
|
||
where
|
||
go (offset : Nat) (m : LevelMap) (i : Nat) (e : Expr) : Expr :=
|
||
if i > 0 then
|
||
let param := params[i-1]!
|
||
let domain := param.type.abstract' (offset - 1) m
|
||
go (offset - 1) m (i - 1) (.lam param.binderName domain e .default)
|
||
else
|
||
e
|
||
|
||
private abbrev _root_.Lean.FVarId.toExprM (fvarId : FVarId) : ToExprM Expr :=
|
||
return fvarId.toExpr (← read) (← get)
|
||
|
||
@[inline] def abstractM (e : Expr) : ToExprM Expr :=
|
||
return e.abstract' (← read) (← get)
|
||
|
||
@[inline] def withFVar (fvarId : FVarId) (k : ToExprM α) : ToExprM α := do
|
||
let offset ← read
|
||
modify fun s => s.insert fvarId offset
|
||
withReader (·+1) k
|
||
|
||
@[inline] partial def withParams (params : Array (Param pu)) (k : ToExprM α) : ToExprM α :=
|
||
go 0
|
||
where
|
||
@[specialize] go (i : Nat) : ToExprM α := do
|
||
if h : i < params.size then
|
||
withFVar params[i].fvarId (go (i+1))
|
||
else
|
||
k
|
||
|
||
@[inline] def run (x : ToExprM α) (offset : Nat := 0) (levelMap : LevelMap := {}) : α :=
|
||
x |>.run offset |>.run' levelMap
|
||
|
||
@[inline] def run' (x : ToExprM α) (xs : Array FVarId) : α :=
|
||
let map := xs.foldl (init := {}) fun map x => map.insert x map.size
|
||
run x xs.size map
|
||
|
||
end ToExpr
|
||
|
||
open ToExpr
|
||
|
||
private def Arg.toExprM (arg : Arg pu) : ToExprM Expr :=
|
||
return arg.toExpr.abstract' (← read) (← get)
|
||
|
||
mutual
|
||
partial def FunDecl.toExprM (decl : FunDecl pu) : ToExprM Expr :=
|
||
withParams decl.params do mkLambdaM decl.params (← decl.value.toExprM)
|
||
|
||
partial def Code.toExprM (code : Code pu) : ToExprM Expr := do
|
||
match code with
|
||
| .let decl k =>
|
||
let type ← abstractM decl.type
|
||
let value ← abstractM decl.value.toExpr
|
||
let body ← withFVar decl.fvarId k.toExprM
|
||
return .letE decl.binderName type value body true
|
||
| .fun decl k _ | .jp decl k =>
|
||
let type ← abstractM decl.type
|
||
let value ← decl.toExprM
|
||
let body ← withFVar decl.fvarId k.toExprM
|
||
return .letE decl.binderName type value body true
|
||
| .return fvarId => fvarId.toExprM
|
||
| .jmp fvarId args => return mkAppN (← fvarId.toExprM) (← args.mapM Arg.toExprM)
|
||
| .unreach type => return mkApp (mkConst ``lcUnreachable) (← abstractM type)
|
||
| .cases c =>
|
||
let alts ← c.alts.mapM fun
|
||
| .alt ctorName params k _ => do
|
||
let body ← withParams params do mkLambdaM params (← k.toExprM)
|
||
return mkApp (mkConst ctorName) body
|
||
| .default k => k.toExprM
|
||
| .ctorAlt i k _ => do
|
||
let body ← k.toExprM
|
||
return mkApp (mkConst i.name) body
|
||
return mkAppN (mkConst `cases) (#[← c.discr.toExprM] ++ alts)
|
||
| .sset fvarId i offset y ty k _ =>
|
||
let value := mkApp5 (mkConst `sset) (.fvar fvarId) (toExpr i) (toExpr offset) (.fvar y) ty
|
||
let body ← withFVar fvarId k.toExprM
|
||
return .letE `dummy (mkConst ``Unit) value body true
|
||
| .oset fvarId offset y k _ =>
|
||
let value := mkApp3 (mkConst `oset) (.fvar fvarId) (toExpr offset) (← y.toExprM)
|
||
let body ← withFVar fvarId k.toExprM
|
||
return .letE `dummy (mkConst ``Unit) value body true
|
||
| .uset fvarId offset y k _ =>
|
||
let value := mkApp3 (mkConst `uset) (.fvar fvarId) (toExpr offset) (.fvar y)
|
||
let body ← withFVar fvarId k.toExprM
|
||
return .letE `dummy (mkConst ``Unit) value body true
|
||
| .setTag fvarId cidx k _ =>
|
||
let body ← withFVar fvarId k.toExprM
|
||
let value := mkApp2 (mkConst `setTag) (.fvar fvarId) (toExpr cidx)
|
||
return .letE `dummy (mkConst ``Unit) value body true
|
||
| .inc fvarId n check persistent k _ =>
|
||
let value := mkApp4 (mkConst `inc) (.fvar fvarId) (toExpr n) (toExpr check) (toExpr persistent)
|
||
let body ← withFVar fvarId k.toExprM
|
||
return .letE `dummy (mkConst ``Unit) value body true
|
||
| .dec fvarId n check persistent k _ =>
|
||
let body ← withFVar fvarId k.toExprM
|
||
let value := mkApp4 (mkConst `dec) (.fvar fvarId) (toExpr n) (toExpr check) (toExpr persistent)
|
||
return .letE `dummy (mkConst ``Unit) value body true
|
||
| .del fvarId k _ =>
|
||
let body ← withFVar fvarId k.toExprM
|
||
let value := mkApp (mkConst `del) (.fvar fvarId)
|
||
return .letE `dummy (mkConst ``Unit) value body true
|
||
end
|
||
|
||
public def Code.toExpr (code : Code pu) (xs : Array FVarId := #[]) : Expr :=
|
||
run' code.toExprM xs
|
||
|
||
public def FunDecl.toExpr (decl : FunDecl pu) (xs : Array FVarId := #[]) : Expr :=
|
||
run' decl.toExprM xs
|
||
|
||
end Lean.Compiler.LCNF
|