Proposed syntax tree data structure for Lean

syntax2.lean

/-   syntax   -/

structure position :=
(line : ℕ) (column : ℕ)

structure range :=
(left : position) (right : position)

structure location :=
(file : string) (range : range)

inductive origin
| file (loc : location)
| synthetic

meta inductive syntax
| num (n : ℕ)
| name (n : name)
| string (s : string)
| quote (e : expr)
| var (i : ℕ) | lconst (n : name)
| node (o : origin) (kind : name)
       (bound : list name) (children : list syntax)

protected meta def syntax.to_format : syntax → format | s :=
format.group $ format.nest 2 $ match s with
| (syntax.num n) := "(num " ++ n ++ ")"
| (syntax.name n) := "(name `" ++ to_fmt n ++ ")"
| (syntax.string n) := "(string " ++ to_fmt n ++ ")"
| (syntax.quote n) := "(quote " ++ to_fmt n ++ ")"
| (syntax.var n) := "(var " ++ to_fmt n ++ ")"
| (syntax.lconst n) := "(lconst `" ++ to_fmt n ++ ")"
| (syntax.node o k b cs) :=
    let cs := cs.foldr (λ c cs, format.line ++ syntax.to_format c ++ cs) format.nil in
    "(node `" ++ to_fmt k ++ " " ++ to_fmt b ++ cs ++ ")"
end

meta instance : has_to_format syntax := ⟨syntax.to_format⟩
meta instance : has_to_string syntax := ⟨to_string ∘ to_fmt⟩
meta instance : has_repr syntax := ⟨to_string ∘ to_fmt⟩

private meta def syntax.lconsts_core : rb_set name → syntax → rb_set name
| s (syntax.lconst n) := s.insert n
| s (syntax.node _ _ _ cs) := cs.foldl syntax.lconsts_core s
| s _ := s

meta def syntax.lconsts (s : syntax) : list name :=
(syntax.lconsts_core mk_rb_set s).to_list

meta def syntax.abstract_core (lcs : list name) : ℕ → syntax → syntax
| off (syntax.lconst n) :=
    match lcs.indexes_of n with
    | (i::_) := syntax.var (i+off)
    | [] := syntax.lconst n
    end
| off (syntax.node o k b cs) :=
    syntax.node o k b (cs.map (syntax.abstract_core (off+b.length)))
| off s := s

meta def syntax.abstract (s : syntax) (lcs : list name) : syntax :=
s.abstract_core lcs 0

meta def syntax.instantiate_core (subst : list syntax) : ℕ → syntax → syntax
| off (syntax.var i) :=
    if i < off then syntax.var i else
    match subst.nth (i - off) with
    | some repl := repl
    | none := syntax.var i
    end
| off (syntax.node o k b cs) :=
    syntax.node o k b (cs.map (syntax.instantiate_core (off + b.length)))
| off s := s

meta def syntax.instantiate (s : syntax) (subst : list syntax) : syntax :=
s.instantiate_core subst 0

meta def syntax.node' (o : origin) (k : name) (bs : list name) (lcs : list name) (cs : list syntax) : syntax :=
syntax.node o k bs (cs.map (λ c, c.abstract lcs))

private meta def syntax.idents_core : rb_set name → syntax → rb_set name
| s (syntax.node _ `ident _ [syntax.name n]) := s.insert n
| s (syntax.node _ _ _ cs) := cs.foldl syntax.lconsts_core s
| s _ := s

meta def syntax.idents (s : syntax) : list name :=
(syntax.idents_core mk_rb_set s).to_list

/-   name resolution   -/

open tactic

@[reducible] meta def resol_ctx := list (name × name)

meta structure resol_state :=
(lctx : resol_ctx := [])
(in_pattern : bool := ff)
(introduced_binders : list name := [])

meta def resol := state_t resol_state tactic
meta instance : monad resol := by delta resol; apply_instance
meta instance : alternative resol := by delta resol; apply_instance
meta instance {α} : has_coe (tactic α) (resol α) := ⟨state_t.lift⟩

meta def resol.pattern {α} (x : resol α) : resol α := do
s₀ ← state_t.read,
state_t.write {s₀ with in_pattern := tt},
a ← x,
state_t.modify $ λ s₁, {s₁ with in_pattern := s₀.in_pattern},
return a

meta def resol.in_pattern : resol bool :=
resol_state.in_pattern <$> state_t.read

meta def resol.new_binder (n : name) : resol name := do
lc ← mk_fresh_name,
state_t.modify $ λ s, {s with introduced_binders := lc :: s.introduced_binders,
                              lctx := (lc, n) :: s.lctx},
return lc

meta def resol.pretty_of_lc (lc : name) : resol name := do
s ← state_t.read,
⟨_, p⟩ ← returnopt (s.lctx.find (λ a, a.fst = lc)),
return p

meta def resol.scoped (r : resol syntax) : resol syntax := do
s₀ ← state_t.read,
state_t.write {s₀ with introduced_binders := []},
(syntax.node o k [] cs) ← r | fail "resol.scopes: not a syntax node",
s₁ ← state_t.read,
let fv_lcs := s₁.introduced_binders,
fvs ← fv_lcs.mmap resol.pretty_of_lc,
state_t.write {s₁ with introduced_binders := s₀.introduced_binders, lctx := s₀.lctx},
return $ syntax.node' o k fvs fv_lcs cs

meta def resolve : syntax → resol syntax
| syn := do
lctx ← resol_state.lctx <$> state_t.read,
match syn with
| (syntax.node o `ident [] [syntax.name n]) :=
    (do resol.in_pattern >>= guardb ∘ bnot,
        some ⟨lc, _⟩ ← return $ lctx.find (λ x, x.snd = n) | failure,
        return $ syntax.lconst lc)
    <|>
    (do d ← get_decl n,
        return $ syntax.node o `const [] [syntax.name d.to_name])
    <|>
    (do resol.in_pattern >>= guardb,
        lc ← resol.new_binder n,
        return $ syntax.lconst lc)
    <|>
    (do s ← state_t.read,
        (fail ("cannot resolve " ++ n.to_string ++ " in " ++ to_string s.lctx) : resol syntax))

| (syntax.node o `lambda [] [pat, bi, ty, body]) := resol.scoped $ do
    ty' ← resolve ty,
    pat' ← resol.pattern $ resolve pat,
    body' ← resolve body,
    return $ syntax.node o `lambda [] [pat', bi, ty', body']

| (syntax.node o `notation [] (syntax.name `plus :: _ :: es)) :=
    return $ syntax.node o `notation [] ([
        syntax.name `plus,
        syntax.node origin.synthetic `notation.choices [] [
            syntax.node origin.synthetic `const [] [syntax.name `has_add.add]
        ]
    ] ++ es)

| (syntax.node o `match_case [] (rhs::lhs)) := resol.scoped $ do
    lhs' ← resol.pattern $ lhs.mmap resolve,
    rhs' ← resolve rhs,
    return $ syntax.node o `match_case [] (rhs'::lhs')

| (syntax.node o k [] cs) :=
    syntax.node o k [] <$> cs.mmap resolve
| (syntax.node o k b cs) := do
    b_lcs ← b.mmap (λ _, mk_fresh_name),
    let cs' := cs.map (λ s, s.instantiate (b_lcs.map syntax.lconst)),
    cs' ← cs'.mmap resolve,
    return $ syntax.node' o k b b_lcs cs'

| s := return s
end

private meta def pick_fresh_vars : list name → list name → list name × list name
| [] ctx := ([], ctx)
| (b::bs) ctx :=
    if b ∈ ctx then
        let b' := match b with
                  | (name.mk_numeral i n) := name.mk_numeral (i+1) n
                  | n := name.mk_numeral 1 n
                  end
        in pick_fresh_vars (b'::bs) ctx
    else
        let (bs', ctx') := pick_fresh_vars bs (b :: ctx) in
        (b :: bs', ctx')

meta def unresolve : list name → syntax → syntax
| ctx (syntax.node o k bs cs) :=
    let (bs, ctx) := pick_fresh_vars bs ctx in
    let syn_bs := bs.map (λ b, syntax.node origin.synthetic `ident [] [syntax.name b]) in
    syntax.node o k [] (cs.map (λ c, unresolve ctx $ syntax.instantiate c syn_bs))
| ctx s := s

/-   elaboration   -/

@[reducible] meta def elab := tactic

meta def elab.expected_type : elab (option expr) := do
tgt ← target,
return $ match tgt with (expr.mvar _ _) := none | _ := some tgt end

meta def parse_binder_info : name → binder_info
| `default := binder_info.default
| _ := binder_info.default

meta def focus_on {α} (goal : expr) (tac : tactic α) : tactic α := do
goals₀ ← get_goals,
set_goals [goal],
res ← tac,
set_goals goals₀,
cleanup,
return res

meta def expand_core : syntax → elab unit | syn := do
match syn with

| (syntax.node o `placeholder [] []) := skip

| (syntax.node o `lambda [b] [pat, syntax.name bi, ty, body]) := do
ty_mvar ← mk_mvar,
body_mvar ← mk_mvar >>= mk_meta_var,
focus_on ty_mvar $ expand_core ty,
focus_on body_mvar $ expand_core body,
exact (expr.lam b (parse_binder_info bi) ty_mvar body_mvar)

| _ := fail $ "cannot expand unknown syntax:\n" ++ repr syn
end

meta def expand (syn : syntax) (expected_type : option expr := none) : tactic expr := do
expected_type ← match expected_type with none := mk_mvar | some t := return t end,
main_mvar ← mk_meta_var expected_type,
focus_on main_mvar $ expand_core syn,
instantiate_mvars main_mvar

/-   examples   -/


/-
match x with
| list.cons x xs := x
| _ := 42
end
-/
meta def match_ex :=
let o l1 c1 l2 c2 := origin.file ⟨"example1.lean", ⟨l1,c1⟩, ⟨l2,c2⟩⟩ in
syntax.node (o 1 0 4 3) `match [] [
    syntax.node (o 1 6 1 7) `ident [] [syntax.name `x],
    syntax.node (o 2 0 2 21) `match_case [] [
        syntax.node (o 2 20 2 21) `ident [] [syntax.name `x],
        syntax.node (o 2 2 2 16) `app [] [
            syntax.node (o 2 2 2 11) `ident [] [syntax.name `list.cons],
            syntax.node (o 2 12 2 13) `ident [] [syntax.name `x],
            syntax.node (o 2 14 2 16) `ident [] [syntax.name `xs]
        ]
    ],
    syntax.node (o 3 2 3 3) `placeholder [] [],
    syntax.node (o 3 7 3 9) `num [] [syntax.num 42]
]


#eval match_ex
#eval do
(resolved, st) ← resolve match_ex {lctx := [(`x122,`x)]},
trace resolved,
trace $ unresolve [`x] resolved

/-
λ x, x + 2
-/
meta def lam_ex :=
let o l1 c1 l2 c2 := origin.file ⟨"example2.lean", ⟨l1,c1⟩, ⟨l2,c2⟩⟩ in
syntax.node (o 1 0 1 10) `lambda [] [
    syntax.node (o 1 2 1 3) `ident [] [syntax.name `x],
    syntax.name `default,
    syntax.node origin.synthetic `placeholder [] [],
    syntax.node (o 1 5 1 10) `notation [] [
        syntax.name `plus,
        syntax.node origin.synthetic `notation.choices [] [],
        syntax.node (o 1 5 1 6) `ident [] [syntax.name `x],
        syntax.node (o 1 9 1 10) `num [] [syntax.num 2]
    ]
]

#eval do
(resolved, st) ← resolve lam_ex {},
trace resolved,
trace $ unresolve [] resolved