!import "arboricx-common.tri" !Local

-- Indexed Arboricx node section reader.
--
-- Node records in the indexed format are just length-prefixed payloads:
--   u32 payloadLength || payload
-- A payload is one of:
--   0x00
--   0x01 || childIndex:u32be
--   0x02 || leftIndex:u32be || rightIndex:u32be
-- Child indices must point strictly backward in the node array.

readNodeRecord = (bs :
  bindResult (readBytes 4 bs)
    (payloadLength afterPayloadLength :
      bindResult (readBytes (u32BEBytesToNat payloadLength) afterPayloadLength)
        (payload afterPayload :
          ok payload afterPayload)))

nodePayloadKind = (nodePayload : bytesHead nodePayload)

nodePayloadHasTag? = (tag nodePayload :
  triage
    false
    (actualTag : byteEq? actualTag tag)
    (_ _ : false)
    (nodePayloadKind nodePayload))

nodePayloadLeaf? = (nodePayload :
  bytesEq? [(0)] nodePayload)

nodePayloadStem? = (nodePayload :
  and?
    (nodePayloadHasTag? nodePayloadStemTag nodePayload)
    (equal? (bytesLength nodePayload) 5))

nodePayloadFork? = (nodePayload :
  and?
    (nodePayloadHasTag? nodePayloadForkTag nodePayload)
    (equal? (bytesLength nodePayload) 9))

nodePayloadValid? = (nodePayload :
  or?
    (nodePayloadLeaf? nodePayload)
    (or?
      (nodePayloadStem? nodePayload)
      (nodePayloadFork? nodePayload)))

nodePayloadStemChildIndex = (nodePayload :
  u32BEBytesToNat (bytesTake 4 (bytesDrop 1 nodePayload)))

nodePayloadForkLeftIndex = (nodePayload :
  u32BEBytesToNat (bytesTake 4 (bytesDrop 1 nodePayload)))

nodePayloadForkRightIndex = (nodePayload :
  u32BEBytesToNat (bytesTake 4 (bytesDrop 5 nodePayload)))

nodeRecordsHaveInvalidPayload? = y (self nodeRecords :
  matchList
    false
    (nodePayload rest :
      or?
        (not? (nodePayloadValid? nodePayload))
        (self rest))
    nodeRecords)

nodePayloadChildIndices = (nodePayload :
  matchBool
    t
    (matchBool
      (pair (nodePayloadStemChildIndex nodePayload) t)
      (pair (nodePayloadForkLeftIndex nodePayload)
        (pair (nodePayloadForkRightIndex nodePayload) t))
      (nodePayloadStem? nodePayload))
    (nodePayloadLeaf? nodePayload))

-- True iff index n names an element before limit in records.
-- For topologically sorted indexed bundles, every child of record i must
-- satisfy childIndex < i, so searching only the prefix [0, i) validates both
-- bounds and acyclicity.
nodeIndexInPrefix? = y (self n records i limit :
  matchBool
    false
    (matchList
      false
      (_ rest :
        matchBool
          true
          (self n rest (succ i) limit)
          (equal? i n))
      records)
    (equal? i limit))

nodeChildIndicesInPrefix? = y (self childIndices records limit :
  matchList
    true
    (childIndex rest :
      matchBool
        (self rest records limit)
        false
        (nodeIndexInPrefix? childIndex records 0 limit))
    childIndices)

nodePayloadIndicesValid? = (nodePayload i records :
  nodeChildIndicesInPrefix?
    (nodePayloadChildIndices nodePayload)
    records
    i)

nodeRecordsValidIndicesFrom? = y (self allRecords remainingRecords i :
  matchList
    true
    (nodePayload rest :
      matchBool
        (self allRecords rest (succ i))
        false
        (nodePayloadIndicesValid? nodePayload i allRecords))
    remainingRecords)

nodeRecordsValidIndices? = (nodeRecords i :
  nodeRecordsValidIndicesFrom? nodeRecords nodeRecords i)

validateNodeRecords = (nodeRecords rest :
  matchBool
    (err errInvalidNodePayload rest)
    (matchBool
      (ok nodeRecords rest)
      (err errMissingNode rest)
      (nodeRecordsValidIndices? nodeRecords 0))
    (nodeRecordsHaveInvalidPayload? nodeRecords))

readNodeRecords_ = y (self bs nodeCount i acc :
  matchBool
    (ok (reverse acc) bs)
    (bindResult (readNodeRecord bs)
      (nodeRecord afterNodeRecord :
        self afterNodeRecord nodeCount (succ i) (pair nodeRecord acc)))
    (equal? i nodeCount))

readNodeRecords = (nodeCount bs :
  readNodeRecords_ bs nodeCount 0 t)

readNodesSection = (bs :
  bindResult (readBytes 8 bs)
    (nodeCount afterNodeCount :
      bindResult (readNodeRecords (u64BEBytesToNat nodeCount) afterNodeCount)
        (nodeRecords afterNodeRecords :
          bindResult (validateNodeRecords nodeRecords afterNodeRecords)
            (validNodeRecords afterValidNodeRecords :
              ok (pair nodeCount validNodeRecords) afterValidNodeRecords))))

readNodesSectionComplete = (bs :
  bindResult (readNodesSection bs)
    (nodesSection afterNodesSection :
      matchBool
        (ok nodesSection afterNodesSection)
        (err errUnexpectedBytes afterNodesSection)
        (bytesNil? afterNodesSection)))

readArboricxNodesSection = (bs :
  bindResult (readArboricxContainer bs)
    (container afterContainer :
      matchPair
        (_ directory :
          bindResult (sectionBytesOrErr arboricxNodesSectionId directory bs afterContainer)
            (nodesBytes _ :
              bindResult (readNodesSectionComplete nodesBytes)
                (nodesSection _ : ok nodesSection afterContainer)))
        container))

nodesSectionCount = (nodesSection :
  matchPair
    (nodeCount _ : nodeCount)
    nodesSection)

nodesSectionRecords = (nodesSection :
  matchPair
    (_ nodeRecords : nodeRecords)
    nodesSection)

nodePayloadToTreeWith = (self nodeRecords nodePayload :
  matchBool
    (ok t t)
    (matchBool
      (bindResult (self (nodePayloadStemChildIndex nodePayload) nodeRecords)
        (child _ : ok (t child) t))
      (bindResult (self (nodePayloadForkLeftIndex nodePayload) nodeRecords)
        (left _ :
          bindResult (self (nodePayloadForkRightIndex nodePayload) nodeRecords)
            (right _ : ok (pair left right) t)))
      (nodePayloadStem? nodePayload))
    (nodePayloadLeaf? nodePayload))

nodeIndexToTree = y (self nodeIndex nodeRecords :
  (nodePayload :
    matchBool
      (nodePayloadToTreeWith self nodeRecords nodePayload)
      (err errMissingNode t)
      (not? (equal? nodePayload t)))
    (nth nodeIndex nodeRecords))

readArboricxTreeFromIndex = (rootIndexBytes bs :
  bindResult (readArboricxNodesSection bs)
    (nodesSection afterContainer :
      bindResult (nodeIndexToTree (u32BEBytesToNat rootIndexBytes) (nodesSectionRecords nodesSection))
        (tree _ : ok tree afterContainer)))

readArboricxExecutableFromIndex = readArboricxTreeFromIndex