plum

module std

import std/Option

import std/Result

import std/List

import std/Map

import std/Number

import std/Str

import std/Buffer

import std/Err

import std/Bool

# A parsed JSON value. Numbers are always stored as `Float` (JSON doesn't

# distinguish integers from floats the way Plum does), and `JsonList`/

# `JsonMap` nest `Json` itself — the same self-referential shape `List[T]`

# already uses internally for `Node[T].next: Option[Node[T]]`.

enum Json =

  | JsonNull

  | JsonBool(Bool)

  | JsonFloat(Float)

  | JsonStr(Str)

  | JsonList(List[Json])

  | JsonMap(Map[Str, Json])

  # Parses `s` as a single JSON document, requiring the whole string (aside

  # from surrounding whitespace) to be consumed. Called as `Json.fromStr(s)`,

  # matching `Int.fromStr`/`Float.fromStr`'s naming.

  fun fromStr(s: Str) -> Result[Json, JsonParseError] =

    parser := JsonParser(src: s, pos: 0)

    parser.skipSpace()

    r := parser.parseValue()

    match r

      Err(e) =>

        return Err(e)

      Ok(v) =>

        parser.skipSpace()

        if parser.pos < parser.src.length()

          return Err(parser.fail("unexpected trailing characters"))

        return Ok(v)

  # Renders `self` back to its JSON text form. A literal `{`/`}` is written

  # doubled (`{{`/`}}`) — a single, un-doubled brace in a Plum string starts

  # or ends a `{expr}` interpolation instead of standing for itself.

  fun toStr(self) -> Str =

    match self

      JsonNull => "null"

      JsonBool(b) => b ? "true" : "false"

      JsonFloat(f) => f.toStr()

      JsonStr(s) => "\"" + escapeJsonStr(s) + "\""

      JsonList(items) => "[" + items.join(",") + "]"

      JsonMap(pairs) => "{{" + jsonMapToStr(pairs) + "}}"

  fun isNull(self) -> Bool =

    match self

      JsonNull => True

      _ => False

  fun asBool(self) -> Option[Bool] =

    match self

      JsonBool(b) => Some(b)

      _ => None

  fun asFloat(self) -> Option[Float] =

    match self

      JsonFloat(f) => Some(f)

      _ => None

  fun asStr(self) -> Option[Str] =

    match self

      JsonStr(s) => Some(s)

      _ => None

  fun asList(self) -> Option[List[Json]] =

    match self

      JsonList(items) => Some(items)

      _ => None

  fun asMap(self) -> Option[Map[Str, Json]] =

    match self

      JsonMap(pairs) => Some(pairs)

      _ => None

# Carries where in the source a parse failure happened, alongside the

# message — implements the shared `Err` trait (`libs/std/err.plum`) rather

# than being a plain `Str`, so a caller that wants the position can get it

# without re-parsing the message.

enum JsonParseError(Err) =

  | JsonParseError(pos: Int, text: Str)

  fun code(self) -> Int =

    1

  fun msg(self) -> Str =

    "json parse error at position {self.pos}: {self.text}"

  # A bare payload-free `None` here can't be disambiguated to one specific

  # `Option` instantiation from this expression alone (see the README's

  # "Known gaps" note on this residual generics limitation — unrelated to the

  # 3 points fixed above) — an explicit `Option[JsonParseError]` return type

  # sidesteps it.

  fun cause(self) -> Option[JsonParseError] =

    None

# A recursive-descent parser over a `Str` source, tracking its read position

# in `pos`. (The original design threaded a `Readable` stream through here,

# but `Readable` is an undefined trait with no methods anywhere in `libs/std`

# — there is nothing to call on it — so this reads directly from a `Str`

# instead, the same way every other stdlib parser, e.g. `Int.fromStr`, does.)

enum JsonParser =

  | JsonParser(src: Str, pos: Int)

  fun fail(self, text: Str) -> JsonParseError =

    JsonParseError(pos: self.pos, text: text)

  # The raw byte at the current position, or 0 (never a valid JSON byte) past

  # the end of `src` — lets every caller test for "end of input" the same way

  # it tests for any other byte, with no separate bounds check.

  fun peek(self) -> Int =

    self.pos >= self.src.length() ? 0 : self.src.byteAt(self.pos)

  fun peekAt(self, offset: Int) -> Int =

    i := self.pos + offset

    i >= self.src.length() ? 0 : self.src.byteAt(i)

  fun advance(self) -> Unit =

    self.pos = self.pos + 1

  fun skipSpace(self) -> Unit =

    while isSpace(self.peek())

      self.advance()

  fun parseValue(self) -> Result[Json, JsonParseError] =

    self.skipSpace()

    c := self.peek()

    if c == 34

      return self.parseString()

    if c == 91

      return self.parseArray()

    if c == 123

      return self.parseObject()

    if c == 116

      return self.parseLiteral("true", JsonBool(True))

    if c == 102

      return self.parseLiteral("false", JsonBool(False))

    if c == 110

      return self.parseLiteral("null", JsonNull)

    if c == 45 || isDigit(c)

      return self.parseNumber()

    return Err(self.fail("unexpected character"))

  fun parseLiteral(self, lit: Str, value: Json) -> Result[Json, JsonParseError] =

    n := lit.length()

    if !matchesAt(self.src, lit, self.pos, n)

      return Err(self.fail("invalid literal, expected '{lit}'"))

    self.pos = self.pos + n

    return Ok(value)

  fun parseNumber(self) -> Result[Json, JsonParseError] =

    start := self.pos

    if self.peek() == 45

      self.advance()

    if !isDigit(self.peek())

      return Err(self.fail("invalid number"))

    while isDigit(self.peek())

      self.advance()

    if self.peek() == 46

      self.advance()

      if !isDigit(self.peek())

        return Err(self.fail("invalid number"))

      while isDigit(self.peek())

        self.advance()

    if self.peek() == 101 || self.peek() == 69

      self.advance()

      if self.peek() == 43 || self.peek() == 45

        self.advance()

      if !isDigit(self.peek())

        return Err(self.fail("invalid number"))

      while isDigit(self.peek())

        self.advance()

    return self.parseNumberText(self.src.sub(start, self.pos))

  # `Float.fromStr` (`libs/std/float.plum`) doesn't understand exponent

  # notation, so a number that has one is split into its mantissa and

  # exponent around the `e`/`E` and recombined here instead.

  fun parseNumberText(self, text: Str) -> Result[Json, JsonParseError] =

    e_index := exponentIndex(text)

    if e_index < 0

      r := parseFloat(text)

      match r

        Ok(f) =>

          return Ok(JsonFloat(f))

        Err(m) =>

          return Err(self.fail(m))

    mantissa := parseFloat(text.sub(0, e_index))

    exponent := parseInt(text.sub(e_index + 1, text.length()))

    match mantissa

      Err(m) =>

        return Err(self.fail(m))

      Ok(mv) =>

        match exponent

          Err(m) =>

            return Err(self.fail(m))

          Ok(ev) =>

            return Ok(JsonFloat(mv * 10.pow(Float(ev))))

  fun parseString(self) -> Result[Json, JsonParseError] =

    self.advance()

    return self.parseStringBody(Buffer())

  fun parseStringBody(self, buf: Buffer) -> Result[Json, JsonParseError] =

    c := self.peek()

    if c == 0

      return Err(self.fail("unterminated string"))

    if c == 34

      self.advance()

      return Ok(JsonStr(buf.toStr()))

    if c == 92

      return self.parseEscape(buf)

    n := self.src.codePointByteLength(self.pos)

    buf.write(self.src.sub(self.pos, self.pos + n))

    self.pos = self.pos + n

    return self.parseStringBody(buf)

  fun parseEscape(self, buf: Buffer) -> Result[Json, JsonParseError] =

    self.advance()

    esc := self.peek()

    if esc == 34

      buf.write("\"")

      self.advance()

      return self.parseStringBody(buf)

    if esc == 92

      buf.write("\\")

      self.advance()

      return self.parseStringBody(buf)

    if esc == 47

      buf.write("/")

      self.advance()

      return self.parseStringBody(buf)

    if esc == 98

      buf.write("\b")

      self.advance()

      return self.parseStringBody(buf)

    if esc == 102

      buf.write("\f")

      self.advance()

      return self.parseStringBody(buf)

    if esc == 110

      buf.write("\n")

      self.advance()

      return self.parseStringBody(buf)

    if esc == 114

      buf.write("\r")

      self.advance()

      return self.parseStringBody(buf)

    if esc == 116

      buf.write("\t")

      self.advance()

      return self.parseStringBody(buf)

    if esc == 117

      self.advance()

      return self.parseUnicodeEscape(buf)

    return Err(self.fail("invalid escape character"))

  fun parseHex4(self) -> Result[Int, JsonParseError] =

    return self.parseHex4From(0, 0)

  fun parseHex4From(self, value: Int, count: Int) -> Result[Int, JsonParseError] =

    if count == 4

      return Ok(value)

    d := hexDigitValue(self.peek())

    if d < 0

      return Err(self.fail("invalid unicode escape"))

    self.advance()

    return self.parseHex4From(value * 16 + d, count + 1)

  fun parseUnicodeEscape(self, buf: Buffer) -> Result[Json, JsonParseError] =

    hr := self.parseHex4()

    match hr

      Err(e) =>

        return Err(e)

      Ok(cp) =>

        return self.parseUnicodeEscapeCont(buf, cp)

  # A high surrogate (`0xD800`..`0xDBFF`) must be followed immediately by a

  # low surrogate (`0xDC00`..`0xDFFF`) `\u` escape; the pair together encodes

  # one codepoint above `0xFFFF`. Any other codepoint is used as-is.

  fun parseUnicodeEscapeCont(self, buf: Buffer, cp: Int) -> Result[Json, JsonParseError] =

    if cp >= 0xD800 && cp <= 0xDBFF && self.peek() == 92 && self.peekAt(1) == 117

      self.advance()

      self.advance()

      hr := self.parseHex4()

      match hr

        Err(e) =>

          return Err(e)

        Ok(low) =>

          if low < 0xDC00 || low > 0xDFFF

            return Err(self.fail("invalid low surrogate"))

          combined := 0x10000 + {{cp - 0xD800} << 10} + {low - 0xDC00}

          buf.writeRune(combined)

          return self.parseStringBody(buf)

    buf.writeRune(cp)

    return self.parseStringBody(buf)

  fun parseArray(self) -> Result[Json, JsonParseError] =

    self.advance()

    self.skipSpace()

    items := List[Json](head: None, tail: None, size: 0)

    if self.peek() == 93

      self.advance()

      return Ok(JsonList(items))

    return self.parseArrayItems(items)

  fun parseArrayItems(self, items: List[Json]) -> Result[Json, JsonParseError] =

    self.skipSpace()

    r := self.parseValue()

    match r

      Err(e) =>

        return Err(e)

      Ok(v) =>

        items.add(v)

        self.skipSpace()

        c := self.peek()

        if c == 44

          self.advance()

          return self.parseArrayItems(items)

        if c == 93

          self.advance()

          return Ok(JsonList(items))

        return Err(self.fail("expected ',' or ']'"))

  fun parseObject(self) -> Result[Json, JsonParseError] =

    self.advance()

    self.skipSpace()

    pairs := Map[Str, Json]()

    if self.peek() == 125

      self.advance()

      return Ok(JsonMap(pairs))

    return self.parseObjectItems(pairs)

  fun parseObjectItems(self, pairs: Map[Str, Json]) -> Result[Json, JsonParseError] =

    self.skipSpace()

    if self.peek() != 34

      return Err(self.fail("expected string key"))

    kr := self.parseString()

    match kr

      Err(e) =>

        return Err(e)

      Ok(keyJson) =>

        return self.parseObjectValue(pairs, jsonStrValue(keyJson))

  fun parseObjectValue(self, pairs: Map[Str, Json], key: Str) -> Result[Json, JsonParseError] =

    self.skipSpace()

    if self.peek() != 58

      return Err(self.fail("expected ':'"))

    self.advance()

    self.skipSpace()

    vr := self.parseValue()

    match vr

      Err(e) =>

        return Err(e)

      Ok(v) =>

        pairs.set(key, v)

        self.skipSpace()

        c := self.peek()

        if c == 44

          self.advance()

          return self.parseObjectItems(pairs)

        if c == 125

          self.advance()

          return Ok(JsonMap(pairs))

        return Err(self.fail("expected ',' or '}}'"))

fun jsonStrValue(j: Json) -> Str =

  match j

    JsonStr(s) => s

    _ => ""

# `Map` has no built-in `toStr`/`join` (unlike `List`, which uses `join` for

# `Json.toStr`'s `JsonList` case above) — a plain free function taking the

# `Buffer` as a captured-by-reference parameter instead of a closure sidesteps

# `libs/std/list.plum`'s `join` note on closures: reassigning a CAPTURED

# variable inside a closure body doesn't persist across that closure's own

# repeated calls, but mutating a heap object's fields THROUGH a captured

# reference (as `appendJsonPair` does to `buf`) does.

fun jsonMapToStr(pairs: Map[Str, Json]) -> Str =

  buf := Buffer()

  pairs.each(|k, v| appendJsonPair(buf, k, v))

  return buf.toStr()

fun appendJsonPair(buf: Buffer, k: Str, v: Json) -> Unit =

  if !buf.isEmpty()

    buf.write(",")

  buf.write("\"")

  buf.write(escapeJsonStr(k))

  buf.write("\":")

  buf.write(v.toStr())

fun exponentIndex(text: Str) -> Int =

  e := text.indexOf("e")

  e >= 0 ? e : text.indexOf("E")

fun escapeJsonStr(s: Str) -> Str =

  escapeJsonFrom(s, 0, s.length())

# Byte-at-a-time except for the pass-through case, which copies a whole

# UTF-8 codepoint at once — copying a lone continuation byte as its own

# "character" (as a naive byte-by-byte pass-through would) would corrupt any

# non-ASCII text, the same hazard documented on `Str`'s own `escapeChar` in

# `libs/std/str.plum`.

fun escapeJsonFrom(s: Str, i: Int, len: Int) -> Str =

  if i >= len

    return ""

  b := s.byteAt(i)

  if b == 34

    return "\\\"" + escapeJsonFrom(s, i + 1, len)

  if b == 92

    return "\\\\" + escapeJsonFrom(s, i + 1, len)

  if b == 10

    return "\\n" + escapeJsonFrom(s, i + 1, len)

  if b == 13

    return "\\r" + escapeJsonFrom(s, i + 1, len)

  if b == 9

    return "\\t" + escapeJsonFrom(s, i + 1, len)

  if b == 8

    return "\\b" + escapeJsonFrom(s, i + 1, len)

  if b == 12

    return "\\f" + escapeJsonFrom(s, i + 1, len)

  if b < 0x20

    return "\\u00" + hexByteStr(b) + escapeJsonFrom(s, i + 1, len)

  n := s.codePointByteLength(i)

  return s.sub(i, i + n) + escapeJsonFrom(s, i + n, len)

fun hexByteStr(b: Int) -> Str =

  hexDigitChar(b / 16) + hexDigitChar(b % 16)

fun hexDigitChar(d: Int) -> Str =

  match d

    0 => "0"

    1 => "1"

    2 => "2"

    3 => "3"

    4 => "4"

    5 => "5"

    6 => "6"

    7 => "7"

    8 => "8"

    9 => "9"

    10 => "a"

    11 => "b"

    12 => "c"

    13 => "d"

    14 => "e"

    15 => "f"

    _ => "0"

fun hexDigitValue(b: Int) -> Int =

  if b >= 48 && b <= 57

    return b - 48

  if b >= 97 && b <= 102

    return b - 97 + 10

  if b >= 65 && b <= 70

    return b - 65 + 10

  return -1

fun isSpace(c: Int) -> Bool =

  c == 32 || c >= 9 && c <= 13

fun isDigit(c: Int) -> Bool =

  c >= 48 && c <= 57

# ---- regression tests ----

fun parseFloatValue(s: Str) -> Float =

  r := Json.fromStr(s)

  match r

    Ok(v) => v.asFloat().unwrapOr(-1.0f)

    Err(_) => -1.0f

test "parses null, true, and false"

  assert Json.fromStr("null").unwrap().isNull()

  assert Json.fromStr("true").unwrap().asBool().unwrap()

  assert !Json.fromStr("false").unwrap().asBool().unwrap()

test "parses integers, decimals, negatives, and exponents"

  assert parseFloatValue("42") == 42.0f

  assert parseFloatValue("-7") == -7.0f

  assert parseFloatValue("3.5") == 3.5f

  assert parseFloatValue("1e2") == 100.0f

  assert parseFloatValue("1.5e-2") == 0.015f

test "parses a plain string, with escapes"

  assert Json.fromStr("\"hello\"").unwrap().asStr().unwrap() == "hello"

  assert Json.fromStr("\"a\\nb\"").unwrap().asStr().unwrap() == "a\nb"

  assert Json.fromStr("\"\\u0041\"").unwrap().asStr().unwrap() == "A"

test "parses an array of mixed values"

  v := Json.fromStr("[1, \"two\", true, null]").unwrap()

  items := v.asList().unwrap()

  assert items.length() == 4

  assert items.get(0).unwrap().asFloat().unwrap() == 1.0f

  assert items.get(1).unwrap().asStr().unwrap() == "two"

  assert items.get(2).unwrap().asBool().unwrap()

  assert items.get(3).unwrap().isNull()

test "parses a nested object"

  v := Json.fromStr("{{\"a\": 1, \"b\": {{\"c\": [1, 2, 3]}}}}").unwrap()

  pairs := v.asMap().unwrap()

  assert pairs.get("a").unwrap().asFloat().unwrap() == 1.0f

  nested := pairs.get("b").unwrap().asMap().unwrap()

  inner := nested.get("c").unwrap().asList().unwrap()

  assert inner.length() == 3

  assert inner.get(2).unwrap().asFloat().unwrap() == 3.0f

test "reports an error on malformed input"

  r := Json.fromStr("{{\"a\": }}")

  assert r.isErr()

test "round-trips a value through toStr and back"

  original := "[1,\"two\",true,null,{{\"k\":3.5}}]"

  v := Json.fromStr(original).unwrap()

  again := Json.fromStr(v.toStr()).unwrap()

  assert again.toStr() == v.toStr()

test "toStr escapes special characters in strings"

  v := JsonStr("a\"b\\c\nd")

  assert v.toStr() == "\"a\\\"b\\\\c\\nd\""