!import "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 : equal? 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)))

nodeU32FromBytes4 = (b0 b1 b2 b3 :
  u32BEBytesToNat
    (pair b0
      (pair b1
        (pair b2
          (pair b3 t)))))

withNodePayloadStemIndex = (nodePayload shortK indexK :
  matchList
    (shortK t)
    (tag r0 :
      matchList
        (shortK t)
        (b0 r1 :
          matchList
            (shortK t)
            (b1 r2 :
              matchList
                (shortK t)
                (b2 r3 :
                  matchList
                    (shortK t)
                    (b3 _ :
                      indexK (nodeU32FromBytes4 b0 b1 b2 b3))
                    r3) r2) r1) r0) nodePayload)

withNodePayloadForkIndices = (nodePayload shortK indicesK :
  matchList
    (shortK t)
    (tag r0 :
      matchList
        (shortK t)
        (l0 r1 :
          matchList
            (shortK t)
            (l1 r2 :
              matchList
                (shortK t)
                (l2 r3 :
                  matchList
                    (shortK t)
                    (l3 r4 :
                      matchList
                        (shortK t)
                        (r0b r5 :
                          matchList
                            (shortK t)
                            (r1b r6 :
                              matchList
                                (shortK t)
                                (r2b r7 :
                                  matchList
                                    (shortK t)
                                    (r3b _ :
                                      indicesK
                                        (nodeU32FromBytes4 l0 l1 l2 l3)
                                        (nodeU32FromBytes4 r0b r1b r2b r3b)) r7) r6) r5) r4) r3) r2) r1) r0) nodePayload)

nodePayloadStemChildIndex = (nodePayload :
  withNodePayloadStemIndex nodePayload (_ : 0) (index : index))

nodePayloadForkLeftIndex = (nodePayload :
  withNodePayloadForkIndices nodePayload (_ : 0) (left right : left))

nodePayloadForkRightIndex = (nodePayload :
  withNodePayloadForkIndices nodePayload (_ : 0) (left right : right))

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

nodePayloadChildIndices = (nodePayload :
  matchList
    t
    (tag rest :
      lazyBool
        (_ :
          withNodePayloadStemIndex
            nodePayload
            (_ : t)
            (childIndex : pair childIndex t))
        (_ :
          lazyBool
            (_ :
              withNodePayloadForkIndices
                nodePayload
                (_ : t)
                (leftIndex rightIndex :
                  pair leftIndex (pair rightIndex t)))
            (_ : t)
            (equal? tag nodePayloadForkTag))
        (equal? tag nodePayloadStemTag))
    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 records n i limit :
  matchList
    false
    (_ rest :
      matchBool
        false
        (matchBool
          true
          (self rest n (succ i) limit)
          (equal? i n))
        (equal? i limit))
    records)

nodeChildIndicesInPrefix? = y (self childIndices records limit :
  matchList
    true
    (childIndex rest :
      matchBool
        (self rest records limit)
        false
        (nodeIndexInPrefix? records childIndex 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)

nodeBuiltTreeIndex = (entry :
  matchPair
    (index _ : index)
    entry)

nodeBuiltTreeValue = (entry :
  matchPair
    (_ tree : tree)
    entry)

nodeTreeByIndex_ = (self builtTrees targetIndex :
  lazyList
    (_ : err errMissingNode t)
    (entry rest :
      lazyBool
        (_ : ok (nodeBuiltTreeValue entry) t)
        (_ : self rest targetIndex)
        (equal? (nodeBuiltTreeIndex entry) targetIndex))
    builtTrees)

nodeTreeByIndex = (builtTrees targetIndex :
  y nodeTreeByIndex_ builtTrees targetIndex)

nodePayloadToTreeFromBuilt = (builtTrees nodePayload :
  matchList
    (err errInvalidNodePayload t)
    (tag rest :
      lazyBool
        (_ : ok t t)
        (_ :
          lazyBool
            (_ :
              withNodePayloadStemIndex
                nodePayload
                (_ : err errInvalidNodePayload t)
                (childIndex :
                  lazyResult
                    (code after : err code after)
                    (child _ : ok (t child) t)
                    (nodeTreeByIndex builtTrees childIndex)))
            (_ :
              lazyBool
                (_ :
                  withNodePayloadForkIndices
                    nodePayload
                    (_ : err errInvalidNodePayload t)
                    (leftIndex rightIndex :
                      lazyResult
                        (code after : err code after)
                        (left _ :
                          lazyResult
                            (code after : err code after)
                            (right _ : ok (pair left right) t)
                            (nodeTreeByIndex builtTrees rightIndex))
                        (nodeTreeByIndex builtTrees leftIndex)))
                (_ : err errInvalidNodePayload t)
                (equal? tag nodePayloadForkTag))
            (equal? tag nodePayloadStemTag))
        (equal? tag 0))
    nodePayload)

nodeBuildState = (targetIndex i builtTrees :
  pair targetIndex (pair i builtTrees))

nodeBuildStateTargetIndex = (state :
  matchPair
    (targetIndex _ : targetIndex)
    state)

nodeBuildStateI = (state :
  matchPair
    (_ rest :
      matchPair
        (i _ : i)
        rest)
    state)

nodeBuildStateBuiltTrees = (state :
  matchPair
    (_ rest :
      matchPair
        (_ builtTrees : builtTrees)
        rest)
    state)

nodeIndexToTree_ = (self remainingRecords state :
  ((nodeIndex :
    ((i :
      ((builtTrees :
        lazyList
          (_ : err errMissingNode t)
          (nodePayload rest :
            lazyResult
              (code after : err code after)
              (tree _ :
                lazyBool
                  (_ : ok tree t)
                  (_ :
                    self
                      rest
                      (nodeBuildState
                        nodeIndex
                        (succ i)
                        (pair (pair i tree) builtTrees)))
                  (equal? i nodeIndex))
              (nodePayloadToTreeFromBuilt builtTrees nodePayload))
          remainingRecords)
       (nodeBuildStateBuiltTrees state)))
     (nodeBuildStateI state)))
   (nodeBuildStateTargetIndex state)))

nodeIndexToTree = (nodeRecords nodeIndex :
  y nodeIndexToTree_ nodeRecords (nodeBuildState nodeIndex 0 t))

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

readArboricxExecutableFromIndex = readArboricxTreeFromIndex