import Mathlib
import Langlib.Grammars.Unrestricted.Definition

/-! # Indexed Grammar Definitions

This file defines indexed grammars (Aho, 1968), their derivation relation, and the induced
language predicate.

An indexed grammar extends context-free grammars by attaching a *stack of flags* (index symbols)
to each nonterminal in a sentential form. Productions can push flags onto, or pop flags from,
these stacks.

## Main declarations

- `IndexedGrammar` — the grammar structure
- `IndexedGrammar.SententialForm` — sentential forms with stacked nonterminals
- `IndexedGrammar.Transforms` — one-step derivation
- `IndexedGrammar.Derives` — reflexive–transitive closure of `Transforms`
- `IndexedGrammar.Language` — the generated language

## References

* A. V. Aho, "Indexed grammars — an extension of context-free grammars", *JACM* 15(4), 1968.
-/

open Relation List

/-- A symbol on the right-hand side of an indexed production.
    A nonterminal can optionally have a flag pushed onto its inherited stack. -/
inductive IRhsSymbol (T : Type) (N : Type) (F : Type) where
  | terminal : T → IRhsSymbol T N F
  | nonterminal : N → Option F → IRhsSymbol T N F
deriving DecidableEq

/-- A production rule of an indexed grammar.

- `lhs` : the nonterminal on the left-hand side.
- `consume` : if `some f`, the rule can only fire when `f` is on top of the nonterminal's
  stack, and pops it; if `none`, no flag is consumed.
- `rhs` : the right-hand side; each nonterminal inherits the (possibly modified) stack of
  the left-hand side, with an optional additional flag push. -/
structure IRule (T : Type) (N : Type) (F : Type) where
  lhs : N
  consume : Option F
  rhs : List (IRhsSymbol T N F)

/-- An indexed grammar that generates words over the alphabet `T`.

- `nt` : type of nonterminals
- `flag` : type of index symbols (flags)
- `initial` : the start nonterminal
- `rules` : the list of production rules -/
structure IndexedGrammar (T : Type) where
  nt : Type
  flag : Type
  initial : nt
  rules : List (IRule T nt flag)

variable {T : Type}

namespace IndexedGrammar

/-- A symbol in a sentential form of an indexed grammar: either a terminal or a nonterminal
    equipped with its current flag stack. -/
inductive ISym (g : IndexedGrammar T) where
  | terminal : T → ISym g
  | indexed : g.nt → List g.flag → ISym g

/-- Expand the right-hand side of a rule, distributing the given stack `σ` to every
    nonterminal (and pushing any specified flag on top). -/
def expandRhs (g : IndexedGrammar T) (rhs : List (IRhsSymbol T g.nt g.flag))
    (σ : List g.flag) : List (ISym g) :=
  rhs.map fun s =>
    match s with
    | IRhsSymbol.terminal t => ISym.terminal t
    | IRhsSymbol.nonterminal n none => ISym.indexed n σ
    | IRhsSymbol.nonterminal n (some f) => ISym.indexed n (f :: σ)

/-- One step of indexed-grammar transformation. -/
def Transforms (g : IndexedGrammar T) (w₁ w₂ : List (ISym g)) : Prop :=
  ∃ (r : IRule T g.nt g.flag) (u v : List (ISym g)) (σ : List g.flag),
    r ∈ g.rules ∧
    (match r.consume with
     | none   => w₁ = u ++ [ISym.indexed r.lhs σ] ++ v
     | some f => w₁ = u ++ [ISym.indexed r.lhs (f :: σ)] ++ v) ∧
    w₂ = u ++ expandRhs g r.rhs σ ++ v

/-- Any number of steps of indexed-grammar transformation. -/
def Derives (g : IndexedGrammar T) : List (ISym g) → List (ISym g) → Prop :=
  ReflTransGen (Transforms g)

/-- A word (list of terminals) is generated by the grammar if it can be derived from the
    initial nonterminal (with empty stack). -/
def Generates (g : IndexedGrammar T) (w : List T) : Prop :=
  Derives g [ISym.indexed g.initial []] (w.map ISym.terminal)

/-- The language generated by the grammar. -/
def Language (g : IndexedGrammar T) : _root_.Language T :=
  setOf (Generates g)

-- -----------------------------------------------------------------------
-- Basic derivation combinators (mirroring the project's CF/CS toolbox)
-- -----------------------------------------------------------------------

theorem deri_self (g : IndexedGrammar T) (w : List (ISym g)) :
    Derives g w w :=
  ReflTransGen.refl

theorem deri_of_tran {g : IndexedGrammar T} {w₁ w₂ : List (ISym g)}
    (h : Transforms g w₁ w₂) : Derives g w₁ w₂ :=
  ReflTransGen.single h

theorem deri_of_deri_deri {g : IndexedGrammar T} {w₁ w₂ w₃ : List (ISym g)}
    (h₁ : Derives g w₁ w₂) (h₂ : Derives g w₂ w₃) : Derives g w₁ w₃ :=
  h₁.trans h₂

theorem deri_of_deri_tran {g : IndexedGrammar T} {w₁ w₂ w₃ : List (ISym g)}
    (h₁ : Derives g w₁ w₂) (h₂ : Transforms g w₂ w₃) : Derives g w₁ w₃ :=
  h₁.tail h₂

theorem deri_of_tran_deri {g : IndexedGrammar T} {w₁ w₂ w₃ : List (ISym g)}
    (h₁ : Transforms g w₁ w₂) (h₂ : Derives g w₂ w₃) : Derives g w₁ w₃ :=
  h₂.head h₁

theorem deri_with_prefix {g : IndexedGrammar T} (p : List (ISym g))
    {w₁ w₂ : List (ISym g)} (h : Derives g w₁ w₂) :
    Derives g (p ++ w₁) (p ++ w₂) := by
  induction h with
  | refl => exact ReflTransGen.refl
  | tail _ ht ih =>
    apply deri_of_deri_tran ih
    obtain ⟨r, u, v, σ, hr, hlhs, hrhs⟩ := ht
    refine ⟨r, p ++ u, v, σ, hr, ?_, by simp_all [List.append_assoc]⟩
    revert hlhs; cases r.consume <;> intro hlhs <;> simp_all [List.append_assoc]

theorem deri_with_suffix {g : IndexedGrammar T} (s : List (ISym g))
    {w₁ w₂ : List (ISym g)} (h : Derives g w₁ w₂) :
    Derives g (w₁ ++ s) (w₂ ++ s) := by
  induction h with
  | refl => exact ReflTransGen.refl
  | tail _ ht ih =>
    apply deri_of_deri_tran ih
    obtain ⟨r, u, v, σ, hr, hlhs, hrhs⟩ := ht
    refine ⟨r, u, v ++ s, σ, hr, ?_, by simp_all [List.append_assoc]⟩
    revert hlhs; cases r.consume <;> intro hlhs <;> simp_all [List.append_assoc]

end IndexedGrammar