lisp parser in lean 4

Lisp.lean

import Lean
import Lean.Parser
import Mathlib
open Lean.Parsec

namespace Lisp
-- Lisp data types
inductive LispVal where
  | Atom (name : String)
  | List (elements : List LispVal)
  | Number (value : Int)
  | String (value : String)
  | Bool (value : Bool)
deriving Repr, BEq



-- Basic Lisp operations
partial def LispVal.eval : LispVal → LispVal
  | Atom name => Atom name  -- For now, just return the atom
  | List [Atom "+", x, y] =>
    match (eval x, eval y) with
    | (Number a, Number b) => Number (a + b)
    | _ => List [Atom "+", x, y]  -- Return unevaluated if not numbers
  | List [Atom "-", x, y] =>
    match (eval x, eval y) with
    | (Number a, Number b) => Number (a - b)
    | _ => List [Atom "-", x, y]
  | List [Atom "*", x, y] =>
    match (eval x, eval y) with
    | (Number a, Number b) => Number (a * b)
    | _ => List [Atom "*", x, y]
  | List [Atom "/", x, y] =>
    match (eval x, eval y) with
    | (Number a, Number b) => Number (a / b)
    | _ => List [Atom "/", x, y]
  | List (Atom "quote" :: xs) => List xs
  | List xs => List (xs.map eval)
  | val => val  -- Return other values unchanged

-- Parser for Lisp expressions
def parseAtom : Lean.Parsec LispVal := do
  let name ← many1Chars (satisfy (fun c => c.isAlpha || c == '+' || c == '-' || c == '*' || c == '/'))
  pure (LispVal.Atom name)

def parseNumber : Lean.Parsec LispVal := do
  let sign ← optional (pchar '-')
  let digits ← many1Chars digit
  let num := (sign.getD ' ').toString ++ digits
  pure (LispVal.Number num.toInt!)

def parseString : Lean.Parsec LispVal := do
  pchar '"'
  let s ← manyChars (satisfy (· != '"'))
  pchar '"'
  pure (LispVal.String s)

def parseBool : Lean.Parsec LispVal := do
  (pstring "#t" >>= fun _ => pure (LispVal.Bool true)) <|>
  (pstring "#f" >>= fun _ => pure (LispVal.Bool false))

partial def parseExpr : Lean.Parsec LispVal := do
  ws
  parseAtom <|> parseNumber <|> parseString <|> parseBool <|> parseList
where
  parseList : Lean.Parsec LispVal := do
    pchar '('
    ws
    let elements ← Lean.Parsec.sepBy parseExpr ws
    ws
    pchar ')'
    pure (LispVal.List elements)

def parseLisp (input : String) : Except String LispVal :=
  match parseExpr input.iter with
  | Lean.Parsec.ParseResult.success _ res => Except.ok res
  | Lean.Parsec.ParseResult.error _ err => Except.error err


-- Lisp syntactic categories
declare_syntax_cat lisp (behavior := symbol)
declare_syntax_cat lisp.op (behavior := symbol)
-- -- Lisp atom syntax
-- syntax ident : lisp

-- -- Lisp number syntax
syntax num : lisp

/-- Lisp string syntax-/
syntax str : lisp

-- Lisp boolean syntax
syntax "tt" : lisp
syntax "ff" : lisp

/--syntax for operators-/
syntax "+" : lisp.op
syntax "-" : lisp.op
syntax "*" : lisp.op
syntax "/" : lisp.op

syntax "(" lisp.op lisp* ")" : lisp
syntax lisp.op : lisp
-- -- Lisp list syntax
syntax "(" lisp* ")" : lisp

-- -- Lisp quote syntax
-- syntax "!!" lisp : lisp

syntax "lisp% " lisp : term

-- Macro for converting Lisp syntax to LispVal
macro_rules
  -- | `(lisp% $x:ident) => `(LispVal.Atom $(Lean.quote (toString x.getId)))
  | `(lisp% $n:num) => `(LispVal.Number $n)
  | `(lisp% $s:str) => `(LispVal.String $s)
  | `(lisp% tt) => `(LispVal.Bool true)
  | `(lisp% ( $op:lisp.op $x:num $y:num)) => `(match ($op).raw with
    | "+" => LispVal.Number ($x + $y)
    | "-" => LispVal.Number ($x - $y)
    | "*" => LispVal.Number ($x * $y)
    | "/" => LispVal.Number ($x / $y)
    | _ => LispVal.Number 0)

namespace MutableNotation
-- xs[i] := xs[2] + 2
-- xs := xs.set i (xs[2] + 2)
syntax:max term "[" term "]" " := " term : term

macro_rules
  | `($xs[$i] := $v) => `(Id.run do
    let mut xs:= $xs
    xs := xs.set $i $v
    pure xs)
end MutableNotation
#eval Array.set
#eval do
 let xs := #[1, 2, 3]
 xs[2] := xs[1]! + 1



  -- | `(lisp% (+ $x $y)) => `(LispVal.List [lisp% (+ $x $y)])
  -- | `(lisp% #f) => `(LispVal.Bool false)
  -- | `(lisp% [[$x $xs*]]) => `(LispVal.List [lisp% $x, $(Lean.quote xs).map (fun x => `(lisp% $x))*])

#eval lisp% (+ 1 3)


end Lisp

namespace Array

/-- Array comprehension notation -/
declare_syntax_cat compClause

/-- for x in xs -/
syntax "for " term " in " term : compClause

/-- if x -/
syntax ifClause := "if " term
/-- special case for x in xs if pred x-/
syntax "for " term " in " term " if " term : compClause
/-- `#[x | for x in xs]` -/
syntax "#[" term " | " compClause,* "]" : term

/-- Semantics for Array comprehension notation.-/
macro_rules
  | `(#[$t | for $x in $xs]) => `(($xs).map (fun $x ↦ $t))
  -- TODO
  | `(#[$t | for $x in $xs if $p]) => `(Id.run do
    let mut out := #[]
    for $x in $xs do
      if $p then out := out.push $t
    return out)
  | `(#[$t | if $x]) => `(if $x then #[ $t ] else #[])
  | `(#[$t | $c, $cs,*]) => `(Array.join #[#[$t | $cs,*] | $c ])

#eval #[x | for x in #[1,2,3] if x > 2]
#eval #[#[x | for x in #[1,2,3] ] | for x in #[1,2,3]]

private def dropWhile (xs : Array a) (p : a → Bool) : Array a :=
  match xs.findIdx? (fun x => !p x) with
  | none => #[]
  | some i => xs[i:]

/-- Compute the sum of all elements in an array. -/
private def sum [Add a] [Zero a] (arr : Array a) : a :=
  arr.foldr Add.add 0

/-- Compute the product of all elements in an array. -/
private def prod [Mul a] [One a] (arr : Array a) : a :=
  arr.foldr Mul.mul 1

/-- Cartesian product of 2 arrays. Example of the list comprehension notation's flexibility. -/
protected def cartProd (xs : Array a) (ys : Array b) : Array (a × b) := #[(x, y) | for x in xs, for y in ys]

/-- Filters a list of values based on a list of booleans. -/
protected def filterBy (xs: Array a) (bs: Array Bool) : Array a := Id.run do
  let mut out := #[]
  for (x, b) in xs.zip bs do
    if b then out := out.push x
  return out

/-- Inserts the `separator` between the elements of the array. TODO this is List.intersperse-/
protected def intersperse (separator : String) (array : Array String) : String := Id.run do
  let mut out := ""
  for _h: i in [:array.size] do
    if i > 0 then out := out ++ separator
    out := out ++ array[i]
  return out


#eval #[1, 2, 3, 4].sum = 10
#eval #[].sum = 0
#eval #[1, 2, 3, 4].prod = 24
#eval #[].prod = 1



#eval #[2 | for _ in [1,2]]
#eval #[x | for (x, _) in [(1,2),(3,4)]]
#eval #[(x, y) | for x in Array.range 5, for y in Array.range 5, if x + y <= 3]
#eval #[#[1],#[1,2]].join = #[1, 1, 2]
#eval #[x| for x in #[1,2,3]] = #[1, 2, 3]
#eval (#[#[2],#[3]]|>.join) = #[2, 3]
#eval #[4 | if 1 < 0] = #[]
#eval #[4 | if 1 < 3] = #[4]
#eval #[(x, y) | for x in Array.range 5, for y in Array.range 5, if x + y <= 3] = #[(0, 0), (0, 1), (0, 2), (0, 3), (1, 0), (1, 1), (1, 2), (2, 0), (2, 1), (3, 0)]

#eval #[1,2,3].filterBy #[true, false, true] = #[1, 3]
#eval #[].dropWhile (fun x => x < 0)

end Array


namespace Lean.HashMap

/--
Checks if all key-value pairs in a `HashMap` satisfy a given predicate.

This function applies the given predicate `f` to each key-value pair in the `HashMap`.
It returns `true` if all pairs satisfy the predicate, and `false` otherwise.
-/
def _root_.Lean.HashMap.all [Hashable K][BEq K][BEq V] (xs: Lean.HashMap K V) (f: K → V → Bool) : Bool :=
  xs.fold (fun acc k v => acc && f k v) true

/--
Checks if any key-value pairs in a `HashMap` satisfy a given predicate.

This function applies the given predicate `f` to each key-value pair in the `HashMap`.
It returns `true` if any pair satisfies the predicate, and `false` otherwise.

-- TODO does this short circuit? make test case
-/
def _root_.Lean.HashMap.any [Hashable K][BEq K][BEq V] (xs: Lean.HashMap K V) (f: K → V → Bool) : Bool :=
  xs.fold (fun acc k v => acc || f k v) false

instance [Hashable K][BEq K][BEq V] : BEq (Lean.HashMap K V) where
  beq xs ys :=
    xs.size == ys.size && xs.all (fun k v => ys.findD k v == v)


/-- Maps both keys and values of a `HashMap` using a given function.

This function applies the given function `f` to each key-value pair in the `HashMap`,
allowing for both keys and values to be transformed. It returns a new `HashMap` with
the transformed key-value pairs.

TODO(alok): is this just an endofunctor? should it return a hashmap and keep the shape?
-/
def map [BEq K] [Hashable K] [BEq K'] [Hashable K'] (f : K → V → (K' × V')) (m : Lean.HashMap K V) : Lean.HashMap K' V' := Id.run do
  let mut result := .empty
  for (k, v) in m do
    let (k', v') := f k v
    result := result.insert k' v'
  return result

/--Display as #{ a ↦ b, c ↦ d }-/
instance [Hashable a] [BEq a] [Repr a] [Repr b]: Repr (Lean.HashMap a b) where
  reprPrec m _ :=
    let entries := m.toArray.map (fun (k, v) => s!"{repr k} ↦ {repr v}")
    "#{" ++ entries.intersperse ", " ++ "}"

instance [ToString a] [ToString b] [BEq a] [Hashable a] : ToString (Lean.HashMap a b) where
  toString m := Id.run do
    let mut out := #[]
    for (k, v) in m do
      out := out.push s!"{k} ↦ {v}"
    "#{" ++ out.intersperse ", " ++ "}"

/-- Maps the values of a `HashMap` using a given function.

This function applies the given function `f` to each value in the `HashMap`,
keeping the keys unchanged.
-/
def mapValues [BEq k] [Hashable k] (f : v → v') (xs : Lean.HashMap k v) : Lean.HashMap k v' :=
  xs.map ((·, f ·))

/-- This function creates a new `HashMap` with a single key-value pair, using the given `k` and `v` as the key and value respectively. -/
def singleton [BEq K] [Hashable K] (k:K) (v : V) : Lean.HashMap K V := Lean.HashMap.empty.insert k v

/-- Syntax category for `HashMap` items separated by the $\maps$ symbol -/
syntax hashMapItem := term " ↦ " term

/-- Syntax for a `HashMap` like `#{1 ↦ 2}` -/
syntax "#{" hashMapItem,* ","? "}" : term

/-- Semantics for `HashMap` notation.-/
macro_rules
  | `(#{}) => `(Lean.HashMap.empty) -- 0
  | `(#{$k ↦ $v}) => `(Lean.HashMap.singleton $k $v) -- 1
  -- `mergeWith` instead of `.insert` is to ensure left to right order for insertion.
  | `(#{$k ↦ $v, $ks,*}) => `(#{$k ↦ $v}.mergeWith (fun _ _ v => v) #{$ks,*}) -- n

#eval (((1:Nat) - (2:Int)) :Int)
-- Example usages
#eval (#{}: Lean.HashMap Int Int)
#eval #{1 ↦ 2}
#eval #{1 ↦ 1, 2 ↦ 2}
#eval #{}.insert 2 2.0
#eval toString #{1 ↦ 1, 2 ↦ 2}
#eval #{1 ↦ 1, 2 ↦ 2}.mapValues (· + 1)

end Lean.HashMap


namespace Option
/-- Unwraps an option, returning the contained value if it is `some`, or a default value if it is `none`. -/
def unwrapOr [Inhabited a] (val: Option a) (default : a := Inhabited.default) : a :=
  val.getD default

#eval (some 3).unwrapOr
#eval none.unwrapOr 2

end Option

namespace List

/-- Construct a new empty list. -/
def empty: List a := []

end List

/-- Local notation for creating a rational number. -/
local notation a "÷" b => Rat.mk' (num := a) (den := b)

instance : Hashable Rat where hash (r: Rat) := hash (hash r.num, hash r.den)

#eval hash (1 ÷ 4) == hash (mkRat 5 20)
#eval (1 ÷ 2) < (3 ÷ 4:Rat)
#eval (1/2) == (3/4)
#eval (1/2) = (3/4)
#eval (1/2) ≥ (3/4)
#eval (1/2) ≤ (3/4)