ai.nim

proc sortObject(tree: var JsonTree, n: NodePos) =
  var pairs: seq[(string, NodePos)]
  for ch0 in sonsReadonly(tree, n):
    let key = firstSon(NodePos ch0.pos).str
    let value = NodePos(ch0.pos + 2)
    pairs.add((key, value))

  pairs.sort do (a, b: (string, NodePos)) -> int:
    cmp(a[0], b[0])

  var i = n.pos + 1
  for key, value in pairs:
    tree.nodes[i] = toNode(opcodeString, getOrIncl(tree.atoms, key))
    tree.nodes[i+1] = tree.nodes[value.pos]
    i += 2

proc rawTest(tree1, tree2: JsonTree, n1, n2: NodePos): bool =
  if n1.kind != n2.kind: return false
  case n1.kind
  of opcodeNull: return true
  of opcodeObject:
    var sortedTree1 = tree1
    var sortedTree2 = tree2
    sortObject(sortedTree1, n1)
    sortObject(sortedTree2, n2)
    let L1 = span(sortedTree1, n1.pos)
    let L2 = span(sortedTree2, n2.pos)
    if L1 != L2: return false
    for i in 0..<L1:
      let c1 = NodePos(i + n1.pos)
      let c2 = NodePos(i + n2.pos)
      if sortedTree1.nodes[c1.pos] != sortedTree2.nodes[c2.pos]:
        return false
  else:
    let L1 = span(tree1, n1.pos)
    let L2 = span(tree2, n2.pos)
    if L1 != L2: return false
    for i in 0..<L1:
      let c1 = NodePos(i + n1.pos)
      let c2 = NodePos(i + n2.pos)
      case c1.kind
      of opcodeInt, opcodeFloat, opcodeString:
        if c1.str != c2.str: return false
      else:
        if tree1.nodes[c1.pos] != tree2.nodes


##


proc compareObjects(tree1, tree2: JsonTree, n1, n2: NodePos): bool =
  for ch0 in sonsReadonly(tree1, n1):
    let key = firstSon(NodePos ch0.pos).str
    let n2val = rawGet(tree2, n2, key)
    if n2val.isNil:
      return false
    if not rawTest(tree1, tree2, NodePos(ch0.pos + 2), n2val):
      return false

  for ch0 in sonsReadonly(tree2, n2):
    let key = firstSon(NodePos ch0.pos).str
    if rawGet(tree1, n1, key).isNil:
      return false

  return true

proc rawTest(tree1, tree2: JsonTree, n1, n2: NodePos): bool =
  if n1.kind != n2.kind: return false
  case n1.kind
  of opcodeNull: return true
  of opcodeObject: return compareObjects(tree1, tree2, n1, n2)
  else:
    let L1 = span(tree1, n1.pos)
    let L2 = span(tree2, n2.pos)
    if L1 != L2: return false
    for i in 0..<L1:
      let c1 = NodePos(i + n1.pos)
      let c2 = NodePos(i + n2.pos)
      case c1.kind
      of opcodeInt, opcodeFloat, opcodeString:
        if c1.str != c2.str: return false
      else:
        if tree1.nodes[c1.pos] != tree2.nodes[c2.pos]: return false
  return true

proc test*(tree: JsonTree; path: JsonPtr, value: JsonTree): bool =
  let n = posFromPtr(tree, path)
  if n.isNil: raisePathError(path.string)
  rawTest(tree, value, n, rootNodeId)

proc `==`*(a, b: JsonTree): bool {.inline.} = rawTest(a, b, rootNodeId, rootNodeId)

###

type
  SortedTree = distinct JsonTree

proc sortTree(tree: JsonTree, n: NodePos): SortedTree =
  var stack = @[n]
  var sortedNodes = newSeq[Node]()
  var sortedAtoms = BiTable[string]()

  while stack.len > 0:
    let curr = stack.pop()
    case curr.kind
    of opcodeObject:
      var pairs: seq[(string, NodePos)]
      for ch0 in sonsReadonly(tree, curr):
        let key = firstSon(NodePos ch0.pos).str
        let value = NodePos(ch0.pos + 2)
        pairs.add((key, value))
        stack.add(value)

      pairs.sort do (a, b: (string, NodePos)) -> int:
        cmp(a[0], b[0])

      for key, value in pairs:
        sortedNodes.add toNode(opcodeKeyValuePair, 2)
        sortedNodes.add toNode(opcodeString, getOrIncl(sortedAtoms, key))
        sortedNodes.add tree.nodes[value.pos]

    of opcodeArray:
      let L = span(tree, curr.pos)
      for i in countdown(L - 1, 0):
        let c = NodePos(curr.pos + i + 1)
        sortedNodes.add tree.nodes[c.pos]
        stack.add(c)

    else:
      sortedNodes.add tree.nodes[curr.pos]

  result = SortedTree(nodes: sortedNodes, atoms: sortedAtoms)

proc rawTest(tree1, tree2: SortedTree, n1, n2: NodePos): bool =
  if n1.kind != n2.kind: return false
  case n1.kind
  of opcodeNull: return true
  of opcodeObject, opcodeArray:
    let L1 = span(JsonTree(tree1), n1.pos)
    let L2 = span(JsonTree(tree2), n2.pos)
    if L1 != L2: return false
    for i in 0..<L1:
      let c1 = NodePos(n1.pos + i + 1)
      let c2 = NodePos(n2.pos + i + 1)
      if not rawTest(tree1, tree2, c1, c2):
        return false
  else:
    if JsonTree(tree1).nodes[n1.pos] != JsonTree(tree2).nodes[n2.pos]:
      return false
  return true

proc test*(tree: JsonTree; path: JsonPtr, value: JsonTree): bool =
  let n = posFromPtr(tree, path)
  if n.isNil: raisePathError(path.string)
  let sortedTree = sortTree(tree, n)
  let sortedValue = sortTree(value, rootNodeId)
  rawTest(sortedTree, sortedValue, rootNodeId, rootNodeId)

proc `==`*(a, b: JsonTree): bool {.inline.} =
  let sortedA = sortTree(a, rootNodeId)
  let sortedB = sortTree(b, rootNodeId)
  rawTest(sortedA, sortedB, rootNodeId, rootNodeId)


proc removeDuplicateKeys(tree: SortedTree): SortedTree =
  var stack = @[rootNodeId]
  var uniqueNodes = newSeq[Node]()
  var uniqueAtoms = BiTable[string]()

  while stack.len > 0:
    let curr = stack.pop()
    case curr.kind
    of opcodeObject:
      var uniqueKeys = initHashSet[string]()
      var uniquePairs: seq[(string, NodePos)]

      for ch0 in sonsReadonly(JsonTree(tree), curr):
        let key = firstSon(NodePos ch0.pos).str
        let value = NodePos(ch0.pos + 2)
        if key notin uniqueKeys:
          uniqueKeys.incl(key)
          uniquePairs.add((key, value))
          stack.add(value)

      for key, value in uniquePairs:
        uniqueNodes.add toNode(opcodeKeyValuePair, 2)
        uniqueNodes.add toNode(opcodeString, getOrIncl(uniqueAtoms, key))
        uniqueNodes.add JsonTree(tree).nodes[value.pos]

    of opcodeArray:
      let L = span(JsonTree(tree), curr.pos)
      for i in countdown(L - 1, 0):
        let c = NodePos(curr.pos + i + 1)
        uniqueNodes.add JsonTree(tree).nodes[c.pos]
        stack.add(c)

    else:
      uniqueNodes.add JsonTree(tree).nodes[curr.pos]

  result = SortedTree(nodes: uniqueNodes, atoms: uniqueAtoms)

proc parseJson(tree: var JsonTree; p: var JsonParser) =
  case p.tok
  of tkString:
    storeAtom(tree, opcodeString, p.a)
    discard getTok(p)
  of tkInt:
    storeAtom(tree, opcodeInt, p.a)
    discard getTok(p)
  of tkFloat:
    storeAtom(tree, opcodeFloat, p.a)
    discard getTok(p)
  of tkTrue:
    tree.nodes.add Node opcodeTrue
    discard getTok(p)
  of tkFalse:
    tree.nodes.add Node opcodeFalse
    discard getTok(p)
  of tkNull:
    tree.nodes.add Node opcodeNull
    discard getTok(p)
  of tkCurlyLe, tkBracketLe:
    var insertPos: seq[PatchPos] = @[]
    var uniqueKeys = initHashSet[string]()
    while true:
      if insertPos.len > 0 and
          kind(NodePos insertPos[^1]) == opcodeObject and p.tok != tkCurlyRi:
        if p.tok != tkString:
          raiseParseErr(p, "string literal as key")
        else:
          let key = p.a
          if key notin uniqueKeys:
            uniqueKeys.incl(key)
            let patchPos = tree.prepare(opcodeKeyValuePair)
            storeAtom(tree, opcodeString, key)
            insertPos.add patchPos
          discard getTok(p)
          eat(p, tkColon)

      template putVal() =
        if insertPos.len > 0 and kind(NodePos insertPos[^1]) == opcodeKeyValuePair:
          tree.patch insertPos.pop()

      case p.tok
      of tkString, tkInt, tkFloat, tkTrue, tkFalse, tkNull:
        parseJson(tree, p)
        putVal()
        if p.tok == tkComma:
          discard getTok(p)
      of tkCurlyLe:
        insertPos.add tree.prepare(opcodeObject)
        uniqueKeys.clear()
        discard getTok(p)
      of tkBracketLe:
        insertPos.add tree.prepare(opcodeArray)
        discard getTok(p)
      of tkCurlyRi:
        if insertPos.len > 0 and kind(NodePos insertPos[^1]) == opcodeObject:
          tree.patch insertPos.pop()
          putVal()
          discard getTok(p)
          if insertPos.len == 0: break
        else:
          raiseParseErr(p, "{")
        if p.tok == tkComma:
          discard getTok(p)
      of tkBracketRi:
        if insertPos.len > 0 and kind(NodePos insertPos[^1]) == opcodeArray:
          tree.patch insertPos.pop()
          putVal()
          discard getTok(p)
          if insertPos.len == 0: break
        else:
          raiseParseErr(p, "{")
        if p.tok == tkComma:
          discard getTok(p)
      else:
        raiseParseErr(p, "{")
  of tkError, tkCurlyRi, tkBracketRi, tkColon, tkComma, tkEof:
    raiseParseErr(p, "{")

proc insertSkippingDuplicates(dest, tree: var JsonTree; destPos, treePos: NodePos) =
  case tree.nodes[treePos.int].kind
  of opcodeObject:
    var keyPositions = initTable[string, NodePos]()
    for n in sonsReadonly(tree, treePos):
      let key = n.firstSon.str
      if key notin keyPositions:
        # Add a new key-value pair
        let patchPos = prepare(dest, dest, NodePos(dest.nodes.len))
        storeAtom(dest, opcodeString, key)
        insertSkippingDuplicates(dest, tree, NodePos(dest.nodes.len), NodePos(n.int + 2))
        patch dest, patchPos
        keyPositions[key] = NodePos(dest.nodes.len - 2)
  of opcodeArray:
    let patchPos = prepare(dest, dest, NodePos(dest.nodes.len))
    for n in sonsReadonly(tree, treePos):
      insertSkippingDuplicates(dest, tree, NodePos(dest.nodes.len), n)
    patch dest, patchPos
  else:
    dest.nodes.add tree.nodes[treePos.int]