Skip to content

Commit

Permalink
initial version of lsimp
Browse files Browse the repository at this point in the history
  • Loading branch information
lecopivo committed May 31, 2024
1 parent fa53125 commit e21a3ec
Show file tree
Hide file tree
Showing 5 changed files with 743 additions and 698 deletions.
152 changes: 152 additions & 0 deletions SciLean/Tactic/GTrans/MetaLCtxM.lean
Original file line number Diff line number Diff line change
@@ -0,0 +1,152 @@
import Lean
import Qq

import SciLean.Lean.Meta.Basic

open Lean Meta

structure _root_.Lean.Meta.ContextCtx where
lctx : LocalContext := {}
localInstances : LocalInstances := #[]
deriving Inhabited

structure _root_.Lean.Meta.ContextCfg where
config : Meta.Config := {}
defEqCtx? : Option DefEqContext := none
synthPendingDepth : Nat := 0
canUnfold? : Option (Meta.Config → ConstantInfo → CoreM Bool) := none


def Lean.Meta.Context.toCtx (ctx : Meta.Context) : ContextCtx where
lctx := ctx.lctx
localInstances := ctx.localInstances

def Lean.Meta.Context.toCfg (cfg : Meta.Context) : ContextCfg where
config := cfg.config
defEqCtx? := cfg.defEqCtx?
synthPendingDepth := cfg.synthPendingDepth
canUnfold? := cfg.canUnfold?


def _root_.Lean.Meta.Context.mkCtxCfg (ctx : ContextCtx) (cfg : ContextCfg) : Meta.Context :=
{ config := cfg.config
lctx := ctx.lctx
localInstances := ctx.localInstances
defEqCtx? := cfg.defEqCtx?
synthPendingDepth := cfg.synthPendingDepth
canUnfold? := cfg.canUnfold? }

-- TODO: change the monad such that we can only add variables to the context and not remove them
-- or completely changes the context
abbrev MetaLCtxM := ReaderT Meta.ContextCfg $ StateT Meta.ContextCtx $ StateRefT Meta.State CoreM


@[always_inline]
instance : Monad MetaLCtxM := let i := inferInstanceAs (Monad MetaLCtxM); { pure := i.pure, bind := i.bind }

instance {α} [Inhabited α] : Inhabited (MetaLCtxM α) where
default := fun _ ctx => do pure (default,ctx)

instance : MonadLCtx MetaLCtxM where
getLCtx := return (← get).lctx

instance : MonadMCtx MetaLCtxM where
getMCtx := return (← getThe Meta.State).mctx
modifyMCtx f := modifyThe Meta.State fun s => { s with mctx := f s.mctx }

instance : MonadEnv MetaLCtxM where
getEnv := return (← getThe Core.State).env
modifyEnv f := do modifyThe Core.State fun s => { s with env := f s.env, cache := {} }; modifyThe Meta.State fun s => { s with cache := {} }

instance : AddMessageContext MetaLCtxM where
addMessageContext := addMessageContextFull

instance : MonadLift MetaM MetaLCtxM where
monadLift x := fun cfg ctx => do let r ← x (.mkCtxCfg ctx cfg); pure (r, ctx)

protected def MetaLCtxM.saveState : MetaLCtxM (SavedState×ContextCtx) :=
return ({ core := (← getThe Core.State), meta := (← get) }, ⟨← getLCtx, ← getLocalInstances⟩)

def MetaLCtxM.restore (b : SavedState) (ctx : ContextCtx) : MetaLCtxM Unit := do
Core.restore b.core
modify fun s => { s with mctx := b.meta.mctx, zetaDeltaFVarIds := b.meta.zetaDeltaFVarIds, postponed := b.meta.postponed }
modifyThe ContextCtx fun _ => ctx

instance : MonadBacktrack (SavedState×ContextCtx) MetaLCtxM where
saveState := MetaLCtxM.saveState
restoreState s := MetaLCtxM.restore s.1 s.2

@[inline] def MetaLCtxM.run (x : MetaLCtxM α) (cfg : ContextCfg := {}) (ctx : ContextCtx := {}) (s : Meta.State := {}) :
CoreM (α × Meta.State) := do
let ((r,_),s) ← x cfg ctx |>.run s
return (r,s)

@[inline] def MetaLCtxM.run' (x : MetaLCtxM α) (cfg : ContextCfg := {}) (ctx : ContextCtx := {}) (s : Meta.State := {}) : CoreM α := do
let ((r,_),_) ← x cfg ctx |>.run s
return r

@[inline] def MetaLCtxM.runInMeta (a : MetaLCtxM α) (k : α → MetaM β) : MetaM β :=
fun ctx => do
let cfg := ctx.toCfg
let ctx := ctx.toCtx
let (a,ctx) ← a cfg ctx
k a (.mkCtxCfg ctx cfg)

instance [MetaEval α] : MetaEval (MetaLCtxM α) :=
fun env opts x _ => MetaEval.eval env opts x.run' true


instance : MonadControl MetaM MetaLCtxM where
stM := fun α => α × ContextCtx
liftWith := fun f => do
let cfg ← readThe ContextCfg
let f' := (f (fun x c s => do
let (x',ctx') ← x cfg ⟨c.lctx,c.localInstances⟩ s
return (x', ctx')))
f'
restoreM := fun x => do let (a, s) ← liftM x; set s; pure a


def lambdaIntro (e : Expr) : MetaLCtxM (Expr × Array Expr) := do
Meta.lambdaTelescope e fun xs e' => do
let ctx ← getThe ContextCtx
fun _ _ => return ((e',xs), ctx)


def letIntro (e : Expr) : MetaLCtxM (Expr × Array Expr) := do
Meta.letTelescope e fun xs e' => do
let ctx ← getThe ContextCtx
fun _ _ => return ((e',xs), ctx)


/-- Adds let declaration into the local context. Returns newly created free variable.
Similar to `withLetDecl` but runs in `MetaLCtxM` instead of `MetaM`. -/
def introLocalDecl (name : Name) (bi : BinderInfo) (type : Expr) : MetaLCtxM Expr := do
let fvarId ← mkFreshFVarId
fun _ ctx =>
let ctx := {ctx with lctx := ctx.lctx.mkLocalDecl fvarId name type bi .default}
let fvar := Expr.fvar fvarId
return (fvar, ctx)


/-- Adds let declaration into the local context. Returns newly created free variable.
Similar to `withLetDecl` but runs in `MetaLCtxM` instead of `MetaM`. -/
def introLetDecl (name : Name) (type val : Expr) : MetaLCtxM Expr := do
let fvarId ← mkFreshFVarId
fun _ ctx =>
let ctx := {ctx with lctx := ctx.lctx.mkLetDecl fvarId name type val (nonDep := false) .default}
let fvar := Expr.fvar fvarId
return (fvar, ctx)


open Lean Meta Qq in
#eval show MetaLCtxM Unit from do
let e := q(fun a b => a + b + 42)

let (b, xs) ← lambdaIntro e

IO.println s!"lambda body: {← ppExpr b}"
IO.println s!"lambda vars: {← liftM <| xs.mapM ppExpr}"
IO.println s!"lambda: {← ppExpr (← mkLambdaFVars xs b)}"
211 changes: 31 additions & 180 deletions SciLean/Tactic/LSimp/Elab.lean
Original file line number Diff line number Diff line change
@@ -1,19 +1,4 @@
/-
Copyright (c) 2020 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
import Lean.Meta.Tactic.Replace
import Lean.Elab.BuiltinNotation
import Lean.Elab.Tactic.Basic
import Lean.Elab.Tactic.ElabTerm
import Lean.Elab.Tactic.Location
import Lean.Elab.Tactic.Config

import Lean.Elab.Tactic.Simp

import SciLean.Tactic.LSimp.Main
import SciLean.Tactic.LSimp.Notation

namespace SciLean.Tactic.LSimp

Expand All @@ -22,168 +7,34 @@ open TSyntax.Compat

open Lean Meta

def traceSimpCall (stx : Syntax) (usedSimps : Simp.UsedSimps) : MetaM Unit := do
let mut stx := stx
if stx[3].isNone then
stx := stx.setArg 3 (mkNullNode #[mkAtom "only"])
let mut args := #[]
let mut localsOrStar := some #[]
let lctx ← getLCtx
let env ← getEnv
for (thm, _) in usedSimps.toArray.qsort (·.2 < ·.2) do
match thm with
| .decl declName => -- global definitions in the environment
if env.contains declName && !simpOnlyBuiltins.contains declName then
args := args.push (← `(Parser.Tactic.simpLemma| $(mkIdent (← unresolveNameGlobal declName)):ident))
| .fvar fvarId => -- local hypotheses in the context
if let some ldecl := lctx.find? fvarId then
localsOrStar := localsOrStar.bind fun locals =>
if !ldecl.userName.isInaccessibleUserName &&
(lctx.findFromUserName? ldecl.userName).get!.fvarId == ldecl.fvarId then
some (locals.push ldecl.userName)
else
none
-- Note: the `if let` can fail for `simp (config := {contextual := true})` when
-- rewriting with a variable that was introduced in a scope. In that case we just ignore.
| .stx _ thmStx => -- simp theorems provided in the local invocation
args := args.push thmStx
| .other _ => -- Ignore "special" simp lemmas such as constructed by `simp_all`.
pure () -- We can't display them anyway.
if let some locals := localsOrStar then
args := args ++ (← locals.mapM fun id => `(Parser.Tactic.simpLemma| $(mkIdent id):ident))
else
args := args.push (← `(Parser.Tactic.simpStar| *))
let argsStx := if args.isEmpty then #[] else #[mkAtom "[", (mkAtom ",").mkSep args, mkAtom "]"]
stx := stx.setArg 4 (mkNullNode argsStx)
logInfoAt stx[0] m!"Try this: {stx}"


/--
`simpLocation ctx discharge? varIdToLemmaId loc`
runs the simplifier at locations specified by `loc`,
using the simp theorems collected in `ctx`
optionally running a discharger specified in `discharge?` on generated subgoals.
Its primary use is as the implementation of the
`simp [...] at ...` and `simp only [...] at ...` syntaxes,
but can also be used by other tactics when a `Syntax` is not available.
For many tactics other than the simplifier,
one should use the `withLocation` tactic combinator
when working with a `location`.
-/
def simpLocation (ctx : Simp.Context) (discharge? : Option Simp.Discharge := none) (loc : Location) : TacticM Simp.UsedSimps := do
match loc with
| Location.targets hyps simplifyTarget =>
withMainContext do
let fvarIds ← getFVarIds hyps
go fvarIds simplifyTarget
| Location.wildcard =>
withMainContext do
go (← (← getMainGoal).getNondepPropHyps) (simplifyTarget := true)
where
go (fvarIdsToSimp : Array FVarId) (simplifyTarget : Bool) : TacticM Simp.UsedSimps := do
let mvarId ← getMainGoal
let (result?, usedSimps) ← SciLean.Tactic.lsimpGoal mvarId ctx (simplifyTarget := simplifyTarget) (discharge? := discharge?) (fvarIdsToSimp := fvarIdsToSimp)
match result? with
| none => replaceMainGoal []
| some (_, mvarId) => replaceMainGoal [mvarId]
return usedSimps


/-
"simp " (config)? (discharger)? ("only ")? ("[" simpLemma,* "]")? (location)?
-/
@[tactic SciLean.Tactic.LSimp.lsimp] def evalSimp : Tactic := fun stx => do
let { ctx, dischargeWrapper } ← withMainContext <| mkSimpContext stx (eraseLocal := false)
let usedSimps ← dischargeWrapper.with fun discharge? =>
simpLocation ctx discharge? (expandOptLocation stx[5])
if tactic.simp.trace.get (← getOptions) then
traceSimpCall stx usedSimps

open Lean Elab Tactic Conv in
@[tactic SciLean.Tactic.LSimp.lsimp_conv] def evalSimpConv : Tactic := fun stx => do
let { ctx, dischargeWrapper } ← withMainContext <| mkSimpContext stx (eraseLocal := false)

let e ← getLhs
let (r, usedSimps) ← dischargeWrapper.with fun discharge? =>
SciLean.Tactic.simp e ctx discharge?

updateLhs r.expr (← r.getProof' e)

if tactic.simp.trace.get (← getOptions) then
traceSimpCall stx usedSimps

def dsimpLocation (ctx : Simp.Context) (loc : Location) : TacticM Unit := do
match loc with
| Location.targets hyps simplifyTarget =>
withMainContext do
let fvarIds ← getFVarIds hyps
go fvarIds simplifyTarget
| Location.wildcard =>
withMainContext do
go (← (← getMainGoal).getNondepPropHyps) (simplifyTarget := true)
where
go (fvarIdsToSimp : Array FVarId) (simplifyTarget : Bool) : TacticM Unit := do
let mvarId ← getMainGoal
let (result?, usedSimps) ← SciLean.Tactic.ldsimpGoal mvarId ctx (simplifyTarget := simplifyTarget) (fvarIdsToSimp := fvarIdsToSimp)
match result? with
| none => replaceMainGoal []
| some mvarId => replaceMainGoal [mvarId]
if tactic.simp.trace.get (← getOptions) then
traceSimpCall (← getRef) usedSimps

@[tactic SciLean.Tactic.LSimp.ldsimp] def evalDSimp : Tactic := fun stx => do
let { ctx, .. } ← withMainContext <| mkSimpContext stx (eraseLocal := false) (kind := .dsimp)
dsimpLocation ctx (expandOptLocation stx[5])


set_option linter.unusedVariables false in
open Lean Elab Tactic Conv in
@[tactic SciLean.Tactic.LSimp.ldsimp_conv] def evalDSimpConv : Tactic := fun stx => do
let { ctx, dischargeWrapper } ← withMainContext <| mkSimpContext stx (eraseLocal := false)

let e ← getLhs
let (e', usedSimps) ← SciLean.Tactic.dsimp e ctx

changeLhs e'

if tactic.simp.trace.get (← getOptions) then
traceSimpCall stx usedSimps


example (f : Nat → Nat) :
(
let a :=
let b := f 10
let x := 5
let g := hold $ λ n => n + b + 10
let cd := (3, g (f 42))
b + cd.1 + g 10 + cd.2 + x
let z := a + 1
(a*a*z, a+a+z)).1
=
sorry
:=
by
-- ldsimp (config := {zeta := false})
lsimp (config := {zeta := false}) only

-- dsimp (config := {zeta := false})
admit

example (f : Nat → Nat) :
(
let a :=
let b := f 10
b + 10
(a*a, a+a)).1
=
sorry
:=
by
ldsimp (config := {zeta := false})
-- lsimp (config := {zeta := false}) only

-- dsimp (config := {zeta := false})
admit

def callLSimpAux (e : Expr) (k : Expr → Expr → Array Expr → MetaM α) : MetaM α := do

let stateRef : IO.Ref Simp.State ← IO.mkRef {}

let mut simprocs : Simp.Simprocs := {}
simprocs ← simprocs.add ``Mathlib.Meta.FunTrans.fun_trans_simproc false
let r := (lsimp e).run
(Simp.mkDefaultMethodsCore #[simprocs]) {config:={zeta:=false,singlePass:=false},simpTheorems:=#[←getSimpTheorems]} stateRef

let a ← r.runInMeta (fun r =>
k r.expr (r.proof?.getD default) r.vars)

return a


def callLSimp (e : Expr) : MetaM Expr := do
callLSimpAux e (fun e _ vars => do
mkLambdaFVars vars e)



open Lean.Parser.Tactic in
syntax (name:=lsimp_conv) "lsimp" /-(config)? (discharger)? (normalizer)?-/ : conv


open Lean Elab Tactic in
@[tactic lsimp_conv] unsafe def lsimpConv : Tactic := fun _ => do
let e ← Conv.getLhs
let e' ← callLSimp e
Conv.updateLhs e' (← mkAppM ``sorryAx #[← mkEq e e', .const ``Bool.false []])
Loading

0 comments on commit e21a3ec

Please sign in to comment.