Lean で数独を実装する例

えびさんによる例

import Lean

/-
# Lean 4のメタプログラミングで数独チェッカーを作った

**数独をプレイするには一番下にスクロールしてください**

元ネタ : https://xuwei-k.hatenablog.com/entry/2024/09/07/184403
元ネタの元ネタ : https://x.com/jalva_dev/status/1832282610233675829
-/

section Rules
/-! まずは数独のルールをコードで表現 -/

-- 1マス。プレースホルダ(`-`)を`none`で表し、`some 0`から`some 8`までで数字を表す
-- (本来許される数字は1~9だけど、内部表現に0~8を使っても特に支障はない)
abbrev Cell := Option (Fin 9)

-- `Fin n`型の値を作るには、自然数と「それが`n`未満の自然数であるという証明」が必要
#check (⟨8, by simp⟩ : Fin 9)
-- 9以上の数で `Fin 9`型を作ることはできない
#check_failure (⟨10, by simp⟩ : Fin 9)

-- 縦・横・ブロックは `Cell` が9個集まったもの
def NineCellsP : List Cell → Prop := (·.length = 9)
-- `{ x // xが満たすべき条件 }`と書くと`Subtype`型になる
-- この型の値を作るには、`List Cell`の値と「その長さが9であるという証明」が必要
abbrev NineCells := { l // NineCellsP l }

-- `NineCells` の各要素に(プレースホルダ以外の)重複がないことを確かめる関数
def NineCells.isValid : NineCells → Prop := fun l =>
  -- プレースホルダを除外
  let l' : List (Fin 9) := l.val.reduceOption
  -- 重複チェックが標準にあって便利(感想)
  l'.Nodup

instance {cs : NineCells} : Decidable cs.isValid := by
  dsimp [NineCells.isValid]; infer_instance

end Rules

section Helpers
/-! 補助関数を定義する -/

-- 9×9=81個の`List Cell`
def AllCells := { l : List Cell // l.length = 81 }

-- あとで使う
instance : Inhabited AllCells where
  default := ⟨List.range 81 |>.map λ _ => none, by simp⟩

-- `AllCells` から `n`番目の行・列・ブロックを切り出す関数

def AllCells.row (l : AllCells) (n : Fin 9) : NineCells :=
  let ⟨n, nH⟩ := n
  let ⟨l, lH⟩ := l
  let l' := l.drop (n * 9) |>.take 9
  have h : NineCellsP l' := by simp [NineCellsP, l', lH]; omega
  ⟨l', h⟩

def AllCells.col (l : AllCells) (n : Fin 9) : NineCells :=
  let ⟨n, nH⟩ := n
  let ⟨l, lH⟩ := l
  let positions := List.range' n 9 9
  let l' : List Cell := positions.attach.map fun ⟨p, h⟩ => by
    simp [positions, List.mem_range'] at h
    refine l.get ⟨p, ?_⟩
    cases h; omega
  ⟨l', by simp [NineCellsP, l', positions]⟩

def AllCells.block (l : AllCells) (n : Fin 9) : NineCells :=
  let ⟨n, nH⟩ := n
  let ⟨l, lH⟩ := l
  let q := n / 3; let r := n % 3
  have qH : q < 3 := (by omega); have rH : r < 3 := by omega
  let positions' := [0,1,2, 9,10,11, 18,19,20]
  let positions := positions'.map (· + (q * 9 * 3) + (r * 3))
  let l' := positions.attach.map fun ⟨p, h⟩ => by
    simp [positions, List.mem_range'] at h
    refine l.get ⟨p, ?_⟩
    obtain ⟨i, h1, h2⟩ := h
    have iH : i ≤ 20 := by simp [positions'] at h1; omega
    omega
  ⟨l', by simp [NineCellsP, l', positions, positions']⟩

def AllCells.isValid (allCells : AllCells) : Prop :=
  let «0..8» : List (Fin 9) := List.range 9
    |>.attach.map (fun ⟨n, h⟩ => ⟨n, by simp_all [List.range_eq_range']⟩)
  let rows := «0..8».map (allCells.row ·)
  let cols := «0..8».map (allCells.col ·)
  let blocks := «0..8».map (allCells.block ·)
  let nineCells : List NineCells := rows ++ cols ++ blocks
  ∀ x, x ∈ nineCells → x.isValid

instance {ac : AllCells} : Decidable ac.isValid := by
  dsimp [AllCells.isValid]
  apply List.decidableBAll

def AllCells.isFilled (ac : AllCells) : Prop := ∀ c, c ∈ ac.val → c.isSome

instance {ac : AllCells} : Decidable ac.isFilled := by
  dsimp [AllCells.isFilled]
  apply List.decidableBAll

def AllCells.isCleared (ac : AllCells) : Prop := ac.isFilled ∧ ac.isValid

instance {ac : AllCells} : Decidable ac.isCleared := by
  dsimp [AllCells.isCleared]; infer_instance

def DupInfo := Fin 3 × Fin 9
deriving Inhabited

def AllCells.getDup {ac : AllCells} (h : ¬ ac.isValid) : DupInfo :=
  let «0..8» : List (Fin 9) := List.range 9
    |>.attach.map (fun ⟨n, h⟩ => ⟨n, by simp_all [List.range_eq_range']⟩)
  let l : List DupInfo := «0..8».bind (λ n =>
    if ¬ (ac.row n |>.isValid) then [(0, n)]
    else if ¬ (ac.col n |>.isValid) then [(1, n)]
    else if ¬ (ac.block n |>.isValid) then [(2, n)]
    else [])
  l.get! 0

def Pos := (Fin 9 × Fin 9) deriving Repr, ToString
def AllCells.get : (ac : AllCells) → (pos : Pos) → Cell
  | ⟨ac, lenH⟩, (r, c) =>
    let pos : Fin ac.length := ⟨r.val * 9 + c.val, by omega⟩
    ac.get pos

def AllCells.set (ac : AllCells) (pos : Pos) (n : Fin 9) (h : ac.get pos = none)
  : AllCells := by
  obtain ⟨ac, lenH⟩ := ac
  obtain ⟨r, c⟩ := pos
  exists ac.set (r.val * 9 + c.val) (some n)
  simpa

def AllCells.holes (ac : AllCells) : Nat := ac.val.count none

theorem List.set_count {α : Type u} [BEq α] [LawfulBEq α] (xs : List α)
  (a b : α) (neq : a ≠ b) (p : Fin xs.length) (h : xs.get p = a)
  : (xs.set p b).count a + 1 = xs.count a := by
  induction xs
  case nil => nomatch p
  case cons hd tl ih =>
    simp at p
    match hp : p with
    | 0 => simp at h; simp [hp, h, neq]
    | ⟨p + 1, ltH⟩ =>
      simp at h ltH ⊢
      have : count a (tl.set p b) + 1 = count a tl := by
        have := ih ⟨p, ltH⟩ (by simpa)
        rw [← this]
      by_cases hdH : a = hd
      . simpa [← hdH]
      . simpa [hdH]

theorem AllCells.set_holes_eq (ac : AllCells) (pos n h)
  : (ac.set pos n h).holes + 1 = ac.holes := by
  obtain ⟨ac, lenH⟩ := ac
  obtain ⟨r, c⟩ := pos
  dsimp [holes, set, get] at h ⊢
  generalize hp : r.val * 9 + c.val = p
  have h : ac[p] = none := by rw [← h]; congr; rw [hp]
  have := List.set_count ac none (some n) (by simp) ⟨p, by rw [lenH]; omega⟩ h
  rw [← this]

-- Pos := (Fin 9 × Fin 9) -- r1c1とかで指定できる
-- AllCells.set Pos Nat (h : ac.get pos = none) → AllCells
-- AllCells.holes が none の数。探索関数はAllCells.holesを減らす
-- ι  ac.get pos
-- ac.holes > (ac.set _ _ _).holes

/-- `Step ac ac'` は ac' がacの空白を1つ埋めたもので、かつac'に重複がないことを表す -/
inductive AllCells.Step : AllCells → AllCells → Prop
  | mk (ac : AllCells) (pos n h) (validH : (ac.set pos n h).isValid)
    : Step ac (ac.set pos n h)

/-- `AllCells.Step` の推移的閉包 -/
inductive AllCells.StepN : AllCells → AllCells → Prop
  | refl (ac : AllCells) : StepN ac ac
  | step (ac ac') (h : Step ac ac') : StepN ac ac'
  | trans (h1 : StepN ac₁ ac₂) (h2 : StepN ac₂ ac₃) : StepN ac₁ ac₃

def AllCells.Solve (ac : AllCells) : Prop :=
  ∃ ac', ac.StepN ac' ∧ ac'.isCleared

theorem AllCells.step_solve {ac ac' : AllCells}
  (stepH : ac.Step ac') (solveH : ac'.Solve) : ac.Solve := by
  have stepN1 := StepN.step _ _ stepH
  obtain ⟨ac'', stepN2, _⟩ := solveH
  have stepNH := StepN.trans stepN1 stepN2
  exists ac''

theorem AllCells.place_num (ac : AllCells) (pos n)
  (h1 := by decide) (h2 : (ac.set pos n h1).isValid := by decide)
  : ac.Step (ac.set pos n h1) := by apply Step.mk; exact h2

end Helpers

section Command
/-! 数独をプレイするためのコマンドを実装する -/

open Lean Elab Term Command

section Cell
  /-! 数値リテラルや文字列リテラルのように、Cell型リテラルを導入する -/

  -- 1マスはプレースホルダか数字
  declare_syntax_cat cell
  syntax "-" : cell
  syntax num : cell

  -- `[c| 5 ]` のように書くとCell型の値として認識されるようにする
  syntax "[c|" cell "]" : term
  elab_rules : term
    | `([c| - ]) => do
      elabTermEnsuringType (← `(none)) (mkConst ``Cell)
    | `([c| $n:num ]) => do
      let nVal : Nat := n.getNat
      if 1 ≤ nVal ∧ nVal ≤ 9 then
        elabTermEnsuringType (← `(some ⟨$n - 1, by simp⟩)) (mkConst ``Cell)
      else
        throwErrorAt n s!"{nVal} は1から9の範囲に収まっていません"

  -- テスト用
  #eval [c| 1 ]
  #eval [c| - ]
  #eval [c| 9 ]
  -- 数字が`0`だとエラーになることを確かめる
  #check_failure [c| 0 ]

end Cell

section SudokuStx

open PrettyPrinter Delaborator

def AllCells.mk (xs : List Cell) (h : xs.length = 81) : AllCells := ⟨xs, h⟩

elab "[sudoku|" cs:cell,* "]" : term => do
  let cs := cs.getElems
  if cs.size ≠ 81 then
    throwError "マスの数が9×9になっていません"
  let cs := cs.map fun c =>
    let s := Syntax.node .none ``«term[c|_]» #[.atom .none "[c|", c, .atom .none "]"]
    @TSyntax.mk `term s
  elabTermEnsuringType (← `(AllCells.mk [$cs,*] (by simp))) (mkConst ``AllCells)

def cell.mkHole : TSyntax `cell :=
  .mk <| .node .none ``«cell-» #[.atom .none "-"]
def cell.mkNum (n : Nat) : TSyntax `cell :=
  .mk <| .node .none ``cell_ #[Syntax.mkNumLit s!"{n}"]

@[app_unexpander AllCells.mk]
def unexpAllCells : Unexpander
  | `(AllCells.mk [$cs,*] $_) =>
    let cs := cs.getElems
    let cs : Array (TSyntax `cell) := cs.map fun c =>
      match c with
      | `(none) => cell.mkHole
      | `(some $n:num) => cell.mkNum (n.getNat + 1)
      | _ => cell.mkNum 11111 -- おかしい時
    `([sudoku| $cs:cell,* ])
  | _ => throw ()

theorem List.sizeOf_eq (xs : List α) : sizeOf xs = xs.length + 1 := by
  induction xs; { simp }; case cons _ _ ih =>
  simp_arith [ih]

def List.toChunks (xs : List α) (n : Nat) : List (List α) :=
  if xs.isEmpty = false ∧ n > 0 then
    xs.take n :: (xs.drop n).toChunks n
  else []
termination_by xs
decreasing_by
  rw [sizeOf_eq, sizeOf_eq]; simp
  rename_i h; simp at h
  have : xs.length > 0 := by apply length_pos.mpr; simp_all
  omega

def formatSudoku' : Syntax → Std.Format
  | `([sudoku| $cs:cell,* ]) =>
    let cells := cs.getElems.map cell |>.toList
    let cellsBy3 := cells.toChunks 3 |>.map three
    let lines := cellsBy3.toChunks 3 |>.map nine
    let linesBy3 := lines.toChunks 3 |>.map line
    let result := linesBy3.drop 1 |>.foldl (fun a b => a ++ sep ++ b) (linesBy3.getD 0 .nil)
    .line ++ .align false ++ .line ++ result
  | _ => .text "formatterにバグがあります"
  where
    cell (stx : TSyntax `cell) : Std.Format :=
      let t : Std.Format := match stx with
      | `(cell| - ) => .text "-"
      | `(cell| $n:num ) => .text s!"{n.getNat}"
      | _ => .text "TSyntax `cell に想定してないパターンがある"
      .align false ++ t ++ .text " "
    three (xs : List Format) : Std.Format :=
      if let [a, b, c] := xs then a ++ b ++ c else .nil
    nine (xs : List Format) : Std.Format :=
      if let [a, b, c] := xs then a ++ .text "| " ++ b ++ .text "| " ++ c else .nil
    sep : Std.Format :=
      .line ++ .align true ++ .text "---------------------" ++ .line
    line (xs : List Format) : Std.Format :=
      if let [a, b, c] := xs then a ++ .line ++ b ++ .line ++ c else .nil

@[formatter «term[sudoku|_,,]»]
def formatSudoku : Formatter := do
  let s ← get
  let stx := s.stxTrav.cur
  Formatter.pushWhitespace (formatSudoku' stx)

end SudokuStx

def OfNatFin2Nat (e : Expr) : Option Nat := do
  let (_, .lit (.natVal n), _) ← e.app3? ``OfNat.ofNat
    | none
  return n

def posExprToPos (e : Expr) : Option (Nat × Nat) := do
  let (_, _, rE, cE) ← e.app4? ``Prod.mk
  let r ← OfNatFin2Nat rE
  let c ← OfNatFin2Nat cE
  return (r, c)

def nExprToN (e : Expr) : Option Nat := do
  OfNatFin2Nat e

def mkLtProof (n N : Nat) : Expr :=
  let n := n + 1
  let reflProof := mkApp (.const ``Nat.le.refl []) (mkNatLit n)
  step n reflProof (N - n)
  where
    step (n : Nat) (e : Expr) : Nat → Expr
    | 0 => e
    | m + 1 =>
      let cst : Expr := .const ``Nat.le.step []
      let nLit : Expr := mkNatLit n
      let mLit : Expr := mkNatLit ((N - 1) - m)
      step n (mkApp3 cst nLit mLit e) m

def mkCell (n : Nat) : Expr :=
  let optSome : Expr := .const ``Option.some [0]
  let finTy : Expr := .const ``Fin []
  mkApp2 optSome (mkApp finTy (mkNatLit 9)) (mkFin9 n)
  where
    mkFin9 (n : Nat) : Expr :=
      let h := mkLtProof n 9
      mkApp3 (.const ``Fin.mk []) (mkNatLit 9) (mkNatLit n) h

def replaceCons (e : Expr) (pos : Nat) (n : Nat) : Option Expr := do
  let (ty, x, xs) ← e.app3? ``List.cons
  let cst : Expr := .const ``List.cons [0]
  match pos with
  | 0 => return mkApp3 cst ty (mkCell n) xs
  | pos + 1 =>
    let xs' ← replaceCons xs pos n
    mkApp3 cst ty x xs'

def replaceAc (e : Expr) : Option Expr := do
  let (ac, pos, n, _) ← e.app4? ``AllCells.set
  let (ac, _) ← ac.app2? ``AllCells.mk
  let n ← nExprToN n
  let (r, c) ← posExprToPos pos
  let ac ← replaceCons ac (r * 9 + c) n
  let rflCst : Expr := .const ``Eq.refl [1]
  let eqProof : Expr := mkApp2 rflCst (.const ``Nat []) (mkNatLit 81)
  return mkApp2 (.const ``AllCells.mk []) ac eqProof

simproc reduceAllCellsSet (AllCells.set _ _ _ _) := fun e => do
  let some ac := replaceAc e
    | return .continue
  let rflCst : Expr := .const ``Eq.refl [1]
  let eqProof : Expr := mkApp2 rflCst (.const ``AllCells []) ac
  return .done (.mk ac eqProof true)

declare_syntax_cat positions
syntax "(" num "," num ")" : positions
elab "place" pos:positions,+ "<=" v:num : tactic => do
  let numCheck (n : TSyntax `num) : Tactic.TacticM PUnit := do
    let nVal := n.getNat
    if ¬ (1 ≤ nVal ∧ nVal ≤ 9) then
      throwErrorAt n "1~9までの数字で指定してください"
  numCheck v

  let f (x : TSyntax `num) : TSyntax `num :=
    Syntax.mkNumLit s!"{x.getNat - 1}"
  let v := f v
  _ ← pos.getElems.mapM fun p => do
    match p with
    | `(positions| ( $r , $c )) =>
      numCheck r; numCheck c

      let r := f r
      let c := f c
      Tactic.evalTactic (← `(tactic|
        apply AllCells.step_solve (AllCells.place_num _ ($r, $c) $v)
      ))
    | _ => throwUnsupportedSyntax
  Tactic.evalTactic (← `(tactic| simp ))

macro "solved" : tactic => `(tactic| refine ⟨_, AllCells.StepN.refl _, by decide⟩ )

set_option pp.deepTerms true
set_option maxHeartbeats 800000

example : AllCells.Solve [sudoku|
  - , 5 , - ,   2 , - , - ,   4 , - , 6,
  - , 6 , 4 ,   1 , - , - ,   - , 7 , 5,
  - , 2 , 8 ,   - , 6 , 5 ,   1 , - , 3,

  8 , - , 6 ,   - , 7 , 2 ,   - , - , -,
  - , - , 3 ,   8 , - , 1 ,   5 , - , -,
  - , 1 , 7 ,   - , 4 , - ,   - , 2 , 8,

  - , - , 5 ,   9 , - , - ,   - , 4 , 1,
  4 , - , - ,   - , 1 , 3 ,   - , - , 9,
  9 , 8 , - ,   - , 5 , - ,   - , 3 , -
] := by
  simp
  place (4,8), (9,3), (1,1) <= 1
  place (2,7), (9,9), (7,5), (8,3) <= 2
  place (7,2) <= 3
  place (3,4) <= 4
  place (6,1), (4,4), (8,8) <= 5
  place (7,1), (5,8) <= 6
  place (8,2), (5,9), (9,4), (7,7) <= 7
  place (1,8), (8,7), (7,6), (2,5) <= 8
  place (1,3), (2,6), (3,8), (5,5), (4,2), (6,7) <= 9
  place (5,1) <= 2
  place (1,5), (4,7), (6,4), (2,1) <= 3
  place (5,2), (4,9), (9,6) <= 4
  place (9,7), (8,4), (6,6) <= 6
  place (1,6), (3,1) <= 7
  solved
  done

/-
`#sudoku` コマンドを実装する
9マスごとの改行などはルールに含めず、各自が適当にスペースや改行を挿入するものとする
あとパースが面倒になるので末尾のカンマを許容しない
-/
syntax (name := sudoku) "#sudoku" cell,* "end" : command
@[command_elab «sudoku»] unsafe def sudokuElab : CommandElab := fun stx => do
  match stx with
  | `(#sudoku $cs:cell,* end) =>
    -- まず`#sudoku`以降に書かれたマスの情報を構文解析、評価して`List Cell`の値を手にいれる
    let cs : List Cell ← Array.toList <$> cs.getElems.mapM fun c => do
      let cellStx ← `([c| $c ])
      -- 構文解析と評価
      liftTermElabM <| evalTerm Cell (mkConst ``Cell) cellStx

    -- 全体のマスの個数が81になっていなければエラーを出す
    if cs.length ≠ 81 then throwError "マスの数が9×9になっていません"
    let allCells : AllCells := if h : cs.length = 81 then ⟨cs, h⟩ else default

    -- 重複チェック
    let «0..8» : List (Fin 9) := List.range 9
      |>.attach.map (fun ⟨n, h⟩ => ⟨n, by simp_all [List.range_eq_range']⟩)
    «0..8».forM fun n => do
      -- 内部表現から表示用の数値に変換
      let nVal : Nat := n.val + 1
      if ¬ (allCells.row n |>.isValid) then
        throwError s!"{nVal}行目に重複があります"
      if ¬ (allCells.col n |>.isValid) then
        throwError s!"{nVal}列目に重複があります"
      if ¬ (allCells.block n |>.isValid) then
        throwError s!"{nVal}番目のブロックに重複があります"

    -- ちゃんと解けてたらメッセージを出す
    if allCells.val.all (·.isSome) then
      logInfo "🎉 おめでとう！クリアです"

  | _ => throwUnsupportedSyntax

end Command

-- これは元ネタの問題だが、これには正解が無い
#sudoku
  - , 3 , - ,   - , - , - ,   - , - , -,
  - , - , - ,   1 , 9 , 5 ,   - , - , -,
  - , - , 8 ,   - , - , - ,   - , 6 , -,

  8 , - , - ,   - , 6 , - ,   - , - , -,
  4 , - , - ,   8 , - , - ,   - , - , 1,
  - , - , - ,   - , 2 , - ,   - , - , -,

  - , 6 , - ,   - , - , - ,   2 , 8 , 9,
  - , - , - ,   4 , 1 , 9 ,   6 , 3 , 5,
  - , - , - ,   - , - , - ,   - , 7 , -
end

-- ちゃんと正解がある問題を作った
#sudoku
  - , 5 , - ,   2 , - , - ,   4 , - , 6,
  - , 6 , 4 ,   1 , - , - ,   - , 7 , 5,
  - , 2 , 8 ,   - , 6 , 5 ,   1 , - , 3,

  8 , - , 6 ,   - , 7 , 2 ,   - , - , -,
  - , - , 3 ,   8 , - , 1 ,   5 , - , -,
  - , 1 , 7 ,   - , 4 , - ,   - , 2 , 8,

  - , - , 5 ,   9 , - , - ,   - , 4 , 1,
  4 , - , - ,   - , 1 , 3 ,   - , - , 9,
  9 , 8 , - ,   - , 5 , - ,   - , 3 , -
end

The issue body contains a large Lean Sudoku example organized into the Rules, Helpers, and Command sections, but it names no repository file or test. First determine where code examples are stored and how they are built, then confirm the intended location and acceptance criteria before starting; the issue does not define what completion should look like.