## TinyChess — the game.
##
## chess.bend says how a piece steps on an empty board. This file is the
## rest: what stands where, what a capture is, when a king is attacked,
## and when a side is mated. Everything here is computable and has no
## holes, so the checker can settle a claim about a position by running
## it rather than by being argued with.
##
## The opening position, black at the top:
##
##       a  b  c  d
##   4   ♟  ♟  ♚  ♟     12 13 14 15
##   3   ·  ·  ·  ·      8  9 10 11
##   2   ·  ·  ·  ·      4  5  6  7
##   1   ♙  ♔  ♙  ♙      0  1  2  3
##
## Both back ranks are full and the kings stand on different files, so
## neither side begins in opposition and every pawn has somewhere to go.
##
## Sizes are the point of every loop below. There are sixteen squares and
## eight pieces, so a scan is sixteen steps, an attack test is sixteen
## scans, and the whole "can white win on the first move" question is a
## few tens of thousands of steps. That is small enough for the checker
## to simply evaluate, which is why the law about it needs no argument.

import Base
import ./chess.bend as C

# ---- the pieces ----------------------------------------------------------

type Piece is Data:
  Empty{}
  WKing{}
  WPawn{}
  WQueen{}
  WKnight{}
  BKing{}
  BPawn{}
  BQueen{}
  BKnight{}

def is_white(p: Piece) -> Bool:
  match p:
    case Empty{}:
      False{}
    case WKing{}:
      True{}
    case WPawn{}:
      True{}
    case WQueen{}:
      True{}
    case WKnight{}:
      True{}
    case BKing{}:
      False{}
    case BPawn{}:
      False{}
    case BQueen{}:
      False{}
    case BKnight{}:
      False{}

def is_empty_piece(p: Piece) -> Bool:
  match p:
    case Empty{}:
      True{}
    case WKing{}:
      False{}
    case WPawn{}:
      False{}
    case WQueen{}:
      False{}
    case WKnight{}:
      False{}
    case BKing{}:
      False{}
    case BPawn{}:
      False{}
    case BQueen{}:
      False{}
    case BKnight{}:
      False{}

def is_king(p: Piece) -> Bool:
  match p:
    case Empty{}:
      False{}
    case WKing{}:
      True{}
    case WPawn{}:
      False{}
    case WQueen{}:
      False{}
    case WKnight{}:
      False{}
    case BKing{}:
      True{}
    case BPawn{}:
      False{}
    case BQueen{}:
      False{}
    case BKnight{}:
      False{}

# whether a piece belongs to the side to move
def mine(+white: Bool, +p: Piece) -> Bool:
  match white:
    case True{}:
      is_white(p)
    case False{}:
      Bool.and(Bool.not(is_empty_piece(p)), Bool.not(is_white(p)))

def theirs(+white: Bool, +p: Piece) -> Bool:
  mine(Bool.not(white), p)

# ---- the board -----------------------------------------------------------

# the sixteen squares, a1 first. Base's List is kind-polymorphic and a
# plain Data field wants a plain Data list, so this one is its own.
type Cells is Data:
  End{}
  Put{head: Piece, tail: Cells}

type Board is Data:
  Board{cells: Cells}

type Step is Data:
  Step{from: Nat, to: Nat}

def Step.from(m: Step) -> Nat:
  match m:
    case Step{from, to}:
      from

def Step.to(m: Step) -> Nat:
  match m:
    case Step{from, to}:
      to

def cell(c: Cells, s: Nat) -> Piece:
  match c:
    case End{}:
      Empty{}
    case Put{head, tail}:
      match s:
        case 0n:
          head
        case 1n+q:
          cell(tail, q)

def put_cell(c: Cells, s: Nat, +p: Piece) -> Cells:
  match c:
    case End{}:
      End{}
    case Put{head, tail}:
      match s:
        case 0n:
          Put{p, tail}
        case 1n+q:
          Put{head, put_cell(tail, q, p)}

# what stands on a square. Off the board reads as empty, which is what
# makes the scans below safe to run past the last square.
def occupant(+b: Board, +s: Nat) -> Piece:
  match b:
    case Board{cells}:
      cell(cells, s)

def place(+b: Board, +s: Nat, +p: Piece) -> Board:
  match b:
    case Board{cells}:
      Board{put_cell(cells, s, p)}

# choosing between two numbers, which the scans need and Bool.pick is not
def pick_nat(c: Bool, +a: Nat, +b: Nat) -> Nat:
  match c:
    case True{}:
      a
    case False{}:
      b

# ---- lines of sight ------------------------------------------------------

# y lies between x and z along one axis: strictly, or all three equal,
# which is the case where the move does not travel on that axis at all
def mid3(+x: Nat, +y: Nat, +z: Nat) -> Bool:
  Bool.or(
    Bool.and(C.nat_lt(x, y), C.nat_lt(y, z)),
    Bool.or(
      Bool.and(C.nat_lt(z, y), C.nat_lt(y, x)),
      Bool.and(C.nat_eq(x, y), C.nat_eq(y, z))))

# s is strictly inside the line from a to b. On a board four wide there
# are at most two such squares, so testing all sixteen is cheaper than
# working out a direction and stepping along it.
def between(+a: Nat, +s: Nat, +b: Nat) -> Bool:
  Bool.and(
    Bool.and(mid3(C.file(a), C.file(s), C.file(b)),
             mid3(C.rank(a), C.rank(s), C.rank(b))),
    Bool.and(Bool.not(C.nat_eq(s, a)), Bool.not(C.nat_eq(s, b))))

def clear_go(+bd: Board, +a: Nat, +b: Nat, n: Nat) -> Bool:
  match n:
    case 0n:
      True{}
    case 1n+k0:
      +k = k0
      blocked = Bool.and(between(a, k, b), Bool.not(is_empty_piece(occupant(bd, k))))
      Bool.and(Bool.not(blocked), clear_go(bd, a, b, k))

# nothing stands between a and b
def clear(+bd: Board, +a: Nat, +b: Nat) -> Bool:
  clear_go(bd, a, b, 16n)

# ---- what a piece threatens ---------------------------------------------

# whether the piece standing on `from` attacks `to`. A pawn attacks only
# where it could take, never where it could push, which is the rule that
# makes pawn structure mean anything.
# the checker will not match on a computed value, so the piece arrives as
# a parameter and the lookup happens in the caller
def attacks_as(p: Piece, +bd: Board, +from: Nat, +to: Nat) -> Bool:
  match p:
    case Empty{}:
      False{}
    case WKing{}:
      C.king_step(from, to)
    case WPawn{}:
      C.pawn_capture(True{}, from, to)
    case WQueen{}:
      Bool.and(C.queen_line(from, to), clear(bd, from, to))
    case WKnight{}:
      C.knight_step(from, to)
    case BKing{}:
      C.king_step(from, to)
    case BPawn{}:
      C.pawn_capture(False{}, from, to)
    case BQueen{}:
      Bool.and(C.queen_line(from, to), clear(bd, from, to))
    case BKnight{}:
      C.knight_step(from, to)

def attacks(+bd: Board, +from: Nat, +to: Nat) -> Bool:
  attacks_as(occupant(bd, from), bd, from, to)

def attacked_go(+bd: Board, +w: Bool, +s: Nat, n: Nat) -> Bool:
  match n:
    case 0n:
      False{}
    case 1n+k0:
      +k = k0
      here = Bool.and(mine(w, occupant(bd, k)), attacks(bd, k, s))
      Bool.or(here, attacked_go(bd, w, s, k))

# whether the given side attacks a square
def attacked(+bd: Board, +w: Bool, +s: Nat) -> Bool:
  attacked_go(bd, w, s, 16n)

def king_go(+bd: Board, +w: Bool, n: Nat) -> Nat:
  match n:
    case 0n:
      16n
    case 1n+k0:
      +k = k0
      here = Bool.and(mine(w, occupant(bd, k)), is_king(occupant(bd, k)))
      pick_nat(here, k, king_go(bd, w, k))

# where the given side's king stands. A board with no king answers 16,
# which is off the board and therefore attacked by nothing.
def king_square(+bd: Board, +w: Bool) -> Nat:
  king_go(bd, w, 16n)

def in_check(+bd: Board, +w: Bool) -> Bool:
  attacked(bd, Bool.not(w), king_square(bd, w))

# ---- moves ---------------------------------------------------------------

def play(+bd: Board, +m: Step) -> Board:
  +f = Step.from(m)
  t = Step.to(m)
  moved = occupant(bd, f)
  place(place(bd, t, moved), f, Empty{})

# the geometry, plus what the position allows: your own piece is not a
# target, a queen needs a clear line, a pawn pushes only onto an empty
# square and takes only onto an occupied one.
def shape_ok(moving: Piece, +target: Piece, +bd: Board, +w: Bool, +f: Nat, +t: Nat) -> Bool:
  match moving:
    case Empty{}:
      False{}
    case WKing{}:
      C.king_step(f, t)
    case WPawn{}:
      Bool.or(
        Bool.and(C.pawn_push(True{}, f, t), is_empty_piece(target)),
        Bool.and(C.pawn_capture(True{}, f, t), theirs(w, target)))
    case WQueen{}:
      Bool.and(C.queen_line(f, t), clear(bd, f, t))
    case WKnight{}:
      C.knight_step(f, t)
    case BKing{}:
      C.king_step(f, t)
    case BPawn{}:
      Bool.or(
        Bool.and(C.pawn_push(False{}, f, t), is_empty_piece(target)),
        Bool.and(C.pawn_capture(False{}, f, t), theirs(w, target)))
    case BQueen{}:
      Bool.and(C.queen_line(f, t), clear(bd, f, t))
    case BKnight{}:
      C.knight_step(f, t)

def pseudo_legal(+bd: Board, +w: Bool, +m: Step) -> Bool:
  +f = Step.from(m)
  +t = Step.to(m)
  +moving = occupant(bd, f)
  +target = occupant(bd, t)
  ok_owner = Bool.and(mine(w, moving), Bool.not(mine(w, target)))
  Bool.and(ok_owner, shape_ok(moving, target, bd, w, f, t))

# and the rule that separates a move from a step: you may not leave your
# own king attacked
def legal(+bd: Board, +w: Bool, +m: Step) -> Bool:
  Bool.and(pseudo_legal(bd, w, m), Bool.not(in_check(play(bd, m), w)))

# ---- mate ----------------------------------------------------------------

def any_to(+bd: Board, +w: Bool, +f: Nat, n: Nat) -> Bool:
  match n:
    case 0n:
      False{}
    case 1n+k0:
      +k = k0
      Bool.or(legal(bd, w, Step{f, k}), any_to(bd, w, f, k))

def any_from(+bd: Board, +w: Bool, n: Nat) -> Bool:
  match n:
    case 0n:
      False{}
    case 1n+k0:
      +k = k0
      Bool.or(any_to(bd, w, k, 16n), any_from(bd, w, k))

def has_legal_move(+bd: Board, +w: Bool) -> Bool:
  any_from(bd, w, 16n)

# in check, and nothing to do about it
def is_checkmate(+bd: Board, +w: Bool) -> Bool:
  Bool.and(in_check(bd, w), Bool.not(has_legal_move(bd, w)))

def is_stalemate(+bd: Board, +w: Bool) -> Bool:
  Bool.and(Bool.not(in_check(bd, w)), Bool.not(has_legal_move(bd, w)))

# ---- can white win at once? ---------------------------------------------

# "there isn't one". Square 16 is off the board, so no real move can
# collide with it, and a failed law prints this against the move it found.
def no_move() -> Step:
  Step{16n, 16n}

def found(+m: Step) -> Bool:
  Bool.not(C.nat_eq(Step.from(m), 16n))

def pick_move(c: Bool, +a: Step, +b: Step) -> Step:
  match c:
    case True{}:
      a
    case False{}:
      b

def wins_to(+bd: Board, +f: Nat, n: Nat) -> Step:
  match n:
    case 0n:
      no_move()
    case 1n+k0:
      +k = k0
      here = Bool.and(legal(bd, True{}, Step{f, k}),
                      is_checkmate(play(bd, Step{f, k}), False{}))
      pick_move(here, Step{f, k}, wins_to(bd, f, k))

def wins_from(+bd: Board, n: Nat) -> Step:
  match n:
    case 0n:
      no_move()
    case 1n+k0:
      +k = k0
      +here = wins_to(bd, k, 16n)
      pick_move(found(here), here, wins_from(bd, k))

# The move that mates, or no_move() when there is none. Returning the move
# rather than a Bool is what lets a broken law name its own counterexample:
# the checker prints what it observed, so the error carries Step{from, to}.
def winning_first_move(+bd: Board) -> Step:
  wins_from(bd, 16n)

# whether white has a first move that mates black on the spot
def white_wins_in_one(+bd: Board) -> Bool:
  found(winning_first_move(bd))

# ---- the positions -------------------------------------------------------

def row(a: Piece, b: Piece, c: Piece, d: Piece, rest: Cells) -> Cells:
  Put{a, Put{b, Put{c, Put{d, rest}}}}

def empty_row(rest: Cells) -> Cells:
  row(Empty{}, Empty{}, Empty{}, Empty{}, rest)

# rank 1 first: white pawn, white king, white pawn, white pawn, then two
# empty ranks, then black's pawn, pawn, king, pawn
def initial() -> Board:
  Board{row(WPawn{}, WKing{}, WPawn{}, WPawn{},
        empty_row(
        empty_row(
        row(BPawn{}, BPawn{}, BKing{}, BPawn{}, End{}))))}

# the same, with the pawn beside white's king traded for a queen
def initial_with_queen() -> Board:
  Board{row(WPawn{}, WKing{}, WQueen{}, WPawn{},
        empty_row(
        empty_row(
        row(BPawn{}, BPawn{}, BKing{}, BPawn{}, End{}))))}

# a queen each, on the square beside each king, so the position stays
# symmetric under turning the board around
def initial_with_queens() -> Board:
  Board{row(WPawn{}, WKing{}, WQueen{}, WPawn{},
        empty_row(
        empty_row(
        row(BPawn{}, BQueen{}, BKing{}, BPawn{}, End{}))))}

# a knight for white alone, on the square the queen took
def initial_with_knight() -> Board:
  Board{row(WPawn{}, WKing{}, WKnight{}, WPawn{},
        empty_row(
        empty_row(
        row(BPawn{}, BPawn{}, BKing{}, BPawn{}, End{}))))}

# and the same two squares, knights instead
def initial_with_knights() -> Board:
  Board{row(WPawn{}, WKing{}, WKnight{}, WPawn{},
        empty_row(
        empty_row(
        row(BPawn{}, BKnight{}, BKing{}, BPawn{}, End{}))))}