l2.lean

import Lean

open Lean
open IO

inductive SExpr
| atom : String → SExpr
| list : List SExpr → SExpr
deriving Repr, Nonempty

instance : ToString SExpr :=
  let rec pp : SExpr → String
  | .atom a => a
  | .list l => "(" ++ String.intercalate " " (l.map pp) ++ ")"
  ⟨pp⟩

structure State where
  input : String
  pos   : String.Pos
deriving Repr

inductive Err where
  | EOF
deriving Repr

instance : ToString Err where
  toString := fun
    | .EOF => "End of File"

abbrev ParserT := StateT State (Except Err)

def peek? : ParserT (Option Char) := do
  let s ← get
  if s.pos < s.input.endPos then
    pure (s.input.get? s.pos)
  else
    pure none

def next : ParserT Unit := do
  modify fun s => { s with pos := s.input.next s.pos }

partial def skipSpaces : ParserT Unit := do
  let some c ← peek? | pure ()
  if c.isWhitespace then
    next
    skipSpaces
  else
    pure ()

partial def parseAtom : ParserT SExpr := do
  let s ← get
  let start := s.pos
  let rec loop : ParserT Unit := do
    let some c ← peek? | pure ()
    if c.isWhitespace || c == '(' || c == ')' then
      pure ()
    else
      next
      loop
  loop
  let s ← get
  let atom := s.input.extract start s.pos
  pure (.atom atom)

mutual
  partial def parseList : ParserT SExpr := do
    let mut elements := #[]
    while true do
      let some c ← peek? | throw .EOF
      if c == ')' then
        next
        break
      else
        let s ← get
        match parseSExpr.run { input := s.input, pos := s.pos } with
        | .ok (sexpr, nextState) =>
          elements := elements.push sexpr
          modify fun _ => nextState
        | .error err => throw err
    pure (.list elements.toList)

  partial def parseSExpr : ParserT SExpr := do
    skipSpaces
    let some c ← peek? | throw .EOF
    if c == '(' then
      next
      parseList
    else
      parseAtom
end

def parse (input : String) : Except Err SExpr :=
  match (parseSExpr.run { input := input, pos := String.Pos.mk 0 }) with
  | .ok (sexpr, _)  => .ok sexpr
  | .error err => .error err

def parseFile (path : System.FilePath) : IO Unit := do
  let contents ← IO.FS.readFile path
  match parse contents with
  | .ok sexpr => IO.println s!"Parsed: {sexpr}"
  | .error err => IO.eprintln s!"Error: {err}"

def main (args : List String) : IO Unit := do
  match args with
  | [file] => parseFile file
  | _ => IO.eprintln "usage: l2 <file>"