import FormalLab.TypeTheory.DependentTypes
import FormalLab.TypeTheory.StructuralRules
import FormalLab.TypeTheory.TypeClasses

/-!
# 第41章：型推論・単一化・双方向型付け

変数 `f` を関数として適用した瞬間、まだ型が分からなくても「`f` の型は、引数の型から結果の型への
関数型でなければならない」という条件は分かります。型推論は、未知の型へメタ変数を割り当て、構文から
等式制約を集め、その制約を同時に満たす代入を求める問題として分解できます。

本章では単相ラムダ計算に自然数と真偽値を加えた小言語を用います。制約生成、occurs check付きの
一階単一化、推論結果への代入を順に実装します。続いて、型を出力する合成判断と既知の型を入力する
検査判断を分けます。最後に、この小さな推論器とLean 4のelaboratorが共有する考え方と、依存型・
型クラス・高階単一化によって異なる範囲を区別します。

## 未知の型を含む単型と制約

単型の文法を次で定めます。

$$
\tau ::= \mathsf{Nat}\mid\mathsf{Bool}\mid\alpha\mid\tau\to\tau.
$$

ここで `α` は対象言語の型変数ではなく、推論中に解く
メタ変数です。制約は二つの単型の等式 `τ ≐ υ` であり、代入 `S` が全制約の両辺を同じ型へ送るとき
`S` を解と呼びます。制約集合を $E$ とすれば、解の条件は
$\forall(\tau\doteq\upsilon)\in E.\;S(\tau)=S(\upsilon)$ です。
-/

namespace FormalLab.TypeTheory.TypeInference

inductive MonoType where
  | nat
  | bool
  | unknown (identifier : Nat)
  | arrow (domain codomain : MonoType)
  deriving DecidableEq, Repr

structure Equation where
  left : MonoType
  right : MonoType
  deriving DecidableEq, Repr

abbrev Substitution := List (Nat × MonoType)

def occurs (identifier : Nat) : MonoType → Bool
  | .nat => false
  | .bool => false
  | .unknown other => identifier == other
  | .arrow domain codomain => occurs identifier domain || occurs identifier codomain

def replace (identifier : Nat) (replacement : MonoType) : MonoType → MonoType
  | .nat => .nat
  | .bool => .bool
  | .unknown other => if identifier = other then replacement else .unknown other
  | .arrow domain codomain =>
      .arrow (replace identifier replacement domain) (replace identifier replacement codomain)

def applySubstitution (substitution : Substitution) (type : MonoType) : MonoType :=
  substitution.foldl (fun current binding => replace binding.1 binding.2 current) type

def satisfies (substitution : Substitution) (equation : Equation) : Prop :=
  applySubstitution substitution equation.left = applySubstitution substitution equation.right

def solves (substitution : Substitution) (equations : List Equation) : Prop :=
  ∀ equation ∈ equations, satisfies substitution equation

example : occurs 0 (.arrow (.unknown 1) (.unknown 0)) = true := rfl
example : occurs 2 (.arrow (.unknown 1) .nat) = false := rfl
example : replace 0 .nat (.arrow (.unknown 0) (.unknown 1)) =
    .arrow .nat (.unknown 1) := rfl

/-!
`occurs α τ` は `α` が `τ` の中に現れるかを調べます。制約 `α ≐ α → Nat` を
`α ↦ α → Nat` で解こうとすると、代入を繰り返すたびに型が大きくなり、有限な単型は得られません。
occurs checkはこの循環を失敗として検出します。equi-recursive型を許す体系なら同じ方程式を別の意味で
扱えるため、この検査はあらゆる型体系に無条件で共通する規則ではありません。

## 単一化は一つの方程式を小さくする

一階単一化の一段は、関数型同士の方程式を二つの部分方程式へ分解するか、メタ変数を有限な型へ
束縛します。異なる型構成子 `Nat` と `Bool` は失敗します。

$$
\begin{aligned}
\{\tau_1\to\tau_2\doteq\upsilon_1\to\upsilon_2\}\cup E
  &\longmapsto
  \{\tau_1\doteq\upsilon_1,\tau_2\doteq\upsilon_2\}\cup E,\\
\{\alpha\doteq\tau\}\cup E
  &\longmapsto [\tau/\alpha]E
  \qquad(\alpha\notin\mathsf{FV}(\tau)).
\end{aligned}
$$

束縛を得たら、残る方程式と既に得た代入の右辺へ直ちに適用します。これにより後続の一段は、古い
メタ変数束縛を再び解釈せずに進めます。
-/

structure UnifyState where
  equations : List Equation
  substitution : Substitution
  deriving Repr

def replaceEquation (identifier : Nat) (replacement : MonoType)
    (equation : Equation) : Equation :=
  ⟨replace identifier replacement equation.left,
    replace identifier replacement equation.right⟩

def bindMeta (identifier : Nat) (type : MonoType) (state : UnifyState) : Option UnifyState :=
  if type = .unknown identifier then
    some state
  else if occurs identifier type then
    none
  else
    some {
      equations := state.equations.map (replaceEquation identifier type)
      substitution :=
        (identifier, type) :: state.substitution.map
          (fun binding => (binding.1, replace identifier type binding.2))
    }

def unifyStep : UnifyState → Option UnifyState
  | ⟨[], substitution⟩ => some ⟨[], substitution⟩
  | ⟨equation :: rest, substitution⟩ =>
      if equation.left = equation.right then
        some ⟨rest, substitution⟩
      else
        match equation.left, equation.right with
        | .unknown identifier, type => bindMeta identifier type ⟨rest, substitution⟩
        | type, .unknown identifier => bindMeta identifier type ⟨rest, substitution⟩
        | .arrow domain codomain, .arrow domain' codomain' =>
            some ⟨⟨domain, domain'⟩ :: ⟨codomain, codomain'⟩ :: rest, substitution⟩
        | _, _ => none

def unify : Nat → UnifyState → Option Substitution
  | _, ⟨[], substitution⟩ => some substitution
  | 0, _ => none
  | fuel + 1, state =>
      match unifyStep state with
      | none => none
      | some state' => unify fuel state'

def applicationEquations : List Equation :=
  [⟨.arrow (.unknown 0) (.unknown 0), .arrow .nat (.unknown 1)⟩]

def applicationSolution : Substitution :=
  [(1, .nat), (0, .nat)]

theorem applicationEquations_unify :
    unify 4 ⟨applicationEquations, []⟩ = some applicationSolution := rfl

theorem applicationSolution_valid : solves applicationSolution applicationEquations := by
  intro equation membership
  simp [applicationEquations] at membership
  subst equation
  rfl

theorem cyclicEquation_isRejected :
    unify 2 ⟨[⟨.unknown 0, .arrow (.unknown 0) .nat⟩], []⟩ = none := rfl

/-!
`unify` の燃料はLeanに停止を明示するための上限です。`none` は不一致、occurs checkの失敗、燃料切れを
区別しません。実用的な実装なら失敗理由をデータ型で分け、方程式の大きさと未解決変数の数から停止性を
証明します。本章で検査した正当性は、代表制約に返した代入が実際に全方程式を満たす
`applicationSolution_valid` です。一般の健全性・完全性・最汎性をこの例から断定しません。

## 構文から制約を生成する

変数、定数、抽象、適用から成る項を考えます。抽象 `λx.t` では引数へ新しいメタ変数を割り当てます。
適用 `f a` では結果用の新しいメタ変数 `β` を作り、`f` の型が `a` の型から `β` への関数型に等しい
という制約を追加します。
-/

inductive RawTerm where
  | variable (index : Nat)
  | natural (value : Nat)
  | boolean (value : Bool)
  | lambda (body : RawTerm)
  | application (function argument : RawTerm)
  deriving Repr

abbrev Context := List MonoType

structure GenerateState where
  nextMeta : Nat
  equations : List Equation
  deriving Repr

def generate : Context → RawTerm → GenerateState → Option (MonoType × GenerateState)
  | context, .variable index, state =>
      context[index]?.map (fun type => (type, state))
  | _, .natural _, state => some (.nat, state)
  | _, .boolean _, state => some (.bool, state)
  | context, .lambda body, state =>
      let parameterType := .unknown state.nextMeta
      let state' := { state with nextMeta := state.nextMeta + 1 }
      match generate (parameterType :: context) body state' with
      | none => none
      | some (bodyType, finalState) => some (.arrow parameterType bodyType, finalState)
  | context, .application function argument, state =>
      match generate context function state with
      | none => none
      | some (functionType, stateAfterFunction) =>
          match generate context argument stateAfterFunction with
          | none => none
          | some (argumentType, stateAfterArgument) =>
              let resultType := .unknown stateAfterArgument.nextMeta
              let equation := ⟨functionType, .arrow argumentType resultType⟩
              some (resultType, {
                nextMeta := stateAfterArgument.nextMeta + 1
                equations := equation :: stateAfterArgument.equations
              })

structure InferenceResult where
  candidate : MonoType
  inferred : MonoType
  constraints : List Equation
  substitution : Substitution
  deriving Repr

def infer (fuel : Nat) (term : RawTerm) : Option InferenceResult :=
  match generate [] term ⟨0, []⟩ with
  | none => none
  | some (candidate, state) =>
      match unify fuel ⟨state.equations, []⟩ with
      | none => none
      | some substitution => some {
          candidate
          inferred := applySubstitution substitution candidate
          constraints := state.equations
          substitution
        }

def identityTerm : RawTerm :=
  .lambda (.variable 0)

def identityAppliedToFive : RawTerm :=
  .application identityTerm (.natural 5)

theorem identityTerm_infers : (infer 8 identityTerm).map (·.inferred) =
    some (.arrow (.unknown 0) (.unknown 0)) := rfl

theorem identityAppliedToFive_infers :
    (infer 8 identityAppliedToFive).map (·.inferred) = some .nat := rfl

def selfApplication : RawTerm :=
  .lambda (.application (.variable 0) (.variable 0))

theorem selfApplication_isRejected : infer 8 selfApplication = none := rfl

/-!
恒等関数には制約がなく、同じメタ変数が引数と結果に残ります。自然数へ適用すると
`α→α ≐ Nat→β` が生成され、単一化により `α` と `β` がともに `Nat` になります。自己適用では
`α ≐ α→β` が生じ、occurs checkで拒否されます。この失敗は項が実行時に必ず失敗するという判定ではなく、
本章の単純型体系では有限な単型を割り当てられないという判定です。

## 双方向型付けは情報の流れを二つに分ける

全ての項から型を推論する代わりに、判断を二方向へ分けられます。

$$
\Gamma\vdash e\Rightarrow A
\quad\text{（$e$ から型 $A$ を合成する）},
\qquad
\Gamma\vdash e\Leftarrow A
\quad\text{（既知の型 $A$ に対して $e$ を検査する）}.
$$

代表的な方向切替は、合成できた型を検査へ渡す規則と、関数型を入力としてラムダを検査する規則です。

$$
\frac{\Gamma\vdash e\Rightarrow A}{\Gamma\vdash e\Leftarrow A}
\qquad
\frac{\Gamma,x:A\vdash e\Leftarrow B}
     {\Gamma\vdash\lambda x.e\Leftarrow A\to B}.
$$

変数と注釈は型を合成できます。注釈のないラムダは引数型を決められないため、期待される関数型に対して
検査します。適用では関数から矢印型を合成し、その始域に対して引数を検査します。
-/

inductive BidirectionalTerm where
  | variable (index : Nat)
  | natural (value : Nat)
  | boolean (value : Bool)
  | lambda (body : BidirectionalTerm)
  | application (function argument : BidirectionalTerm)
  | annotation (term : BidirectionalTerm) (type : MonoType)
  deriving Repr

mutual
  inductive Synthesizes : Context → BidirectionalTerm → MonoType → Prop where
    | variable : context[index]? = some type →
        Synthesizes context (.variable index) type
    | natural : Synthesizes context (.natural value) .nat
    | boolean : Synthesizes context (.boolean value) .bool
    | application : Synthesizes context function (.arrow domain codomain) →
        Checks context argument domain →
        Synthesizes context (.application function argument) codomain
    | annotation : Checks context term type →
        Synthesizes context (.annotation term type) type

  inductive Checks : Context → BidirectionalTerm → MonoType → Prop where
    | fromSynth : Synthesizes context term type → Checks context term type
    | lambda : Checks (domain :: context) body codomain →
        Checks context (.lambda body) (.arrow domain codomain)
end

def bidirectionalIdentity : BidirectionalTerm :=
  .lambda (.variable 0)

theorem bidirectionalIdentity_checks :
    Checks [] bidirectionalIdentity (.arrow .nat .nat) := by
  exact .lambda (.fromSynth (.variable rfl))

theorem annotatedIdentity_synthesizes :
    Synthesizes [] (.annotation bidirectionalIdentity (.arrow .bool .bool))
      (.arrow .bool .bool) := by
  exact .annotation (.lambda (.fromSynth (.variable rfl)))

/-!
同じラムダが `Nat→Nat` と `Bool→Bool` の両方に検査できますが、注釈なしでは一つの型を合成しません。
注釈は情報の向きを反転させ、検査できた型を外側へ合成します。これは前半の全域的な制約生成とは
異なる設計です。両者を組み合わせる言語では、どこでメタ変数を作り、どこで注釈を要求するかが
推論可能性とエラーメッセージを左右します。

Leanのelaboratorも期待型、メタ変数、制約、型クラス探索を用います。ただし依存型では型が項を含み、
暗黙引数の補完や高階単一化も関わります。本章の `infer` はLean 4のelaboratorの模型であって、その
仕様や完全なアルゴリズムではありません。elaboratorが完全な項を構成した後、kernelが別に型検査します。

## 要点

* 型推論は未知型の導入、制約生成、単一化、代入適用へ分解できる。
* occurs checkは有限単型の中で循環代入を拒み、再帰型を許す体系では前提が変わる。
* 一階単一化は型構成子を分解し、全制約を満たす代入を求める。
* 双方向型付けは型を合成する判断と、既知の型に対して検査する判断を分ける。
* Leanのelaborationはkernelの型検査ではなく、本章の単相推論器より広い問題を扱う。

## 研究史と文献案内

Robinson [ROB65] は自動推論のresolution原理の中で単一化アルゴリズムを提示しました。Hindley [HIN69] は
組合せ論理の主型を扱い、Milner [MIL78] はプログラミング言語MLの文脈で多相型規律とAlgorithm Wを
与えました。Damas–Milner [DM82] は関数型プログラムのprincipal type-schemeを研究します。これらを
一つの同時成立した「Hindley–Milnerアルゴリズム」へ潰さず、問題設定と体系を年代順に読み分けます。

双方向型付けの高階多相への現代的展開はDunfield–Krishnaswami [DK13] を参照してください。本章の単相判断は
同論文の高階ランク体系そのものではありません。Lean 4のメタ変数、単一化、型クラス探索、elaborationの
現行仕様は [LEAN-REF] が対象です。

## 問題

### 制約生成を導出木と実行結果で照合する

項 `(λf. λx. f x) (λy. y)` を `RawTerm` で作り、各部分項へ割り当てられるメタ変数を生成順に記録して
ください。適用ごとに生じる方程式を手で列挙し、`generate` の結果と比較します。十分な燃料で `infer` を
実行し、得られた型へ代入を適用し直して同じ型になることを確認すれば完了です。

### 単一化の不変条件と失敗を分類する

関数型の分解、左右のメタ変数束縛、構成子不一致、occurs check失敗を一例ずつ作ってください。各一段の
前後で、解集合が保存されるか、空になるかを `satisfies` で調べます。燃料切れによる `none` と論理的な
解なしを区別する結果型を設計し、呼出側が報告できる情報の差を説明すれば完了です。

### 合成と検査の境目を設計する

注釈なしラムダ、注釈付きラムダ、二重適用を `BidirectionalTerm` で表してください。各節点を `⇒` と
`⇐` のどちらで読むかを導出木へ明示し、方向を切り替える規則を特定します。ラムダにも常に型を合成させる
設計と比較し、必要なメタ変数、要求する注釈、principal typeの有無がどう変わるかを述べれば完了です。
-/

end FormalLab.TypeTheory.TypeInference