!import "base.tri" !Local
!import "list.tri" !Local
!import "bytes.tri" !Local
!import "binary.tri" !Local

arborixMagic = [(65) (82) (66) (79) (82) (73) (88) (0)]
arborixMajorVersion = [(0) (1)]
arborixMinorVersion = [(0) (0)]
arborixManifestSectionId = [(0) (0) (0) (1)]
arborixNodesSectionId = [(0) (0) (0) (2)]

errMissingSection = 4
errUnsupportedVersion = 5
errDuplicateSection = 6
errDuplicateNode = 7
errInvalidNodePayload = 8
errMissingNode = 9

nodePayloadLeafTag = 0
nodePayloadStemTag = 1
nodePayloadForkTag = 2

readArborixMagic = (bs : expectBytes arborixMagic bs)

readArborixHeader = (bs :
  bindResult (readArborixMagic bs)
    (_ afterMagic :
      bindResult (readBytes 2 afterMagic)
        (majorVersion afterMajor :
          bindResult (readBytes 2 afterMajor)
            (minorVersion afterMinor :
              bindResult (readBytes 4 afterMinor)
                (sectionCount afterSectionCount :
                  bindResult (readBytes 8 afterSectionCount)
                    (flags afterFlags :
                      bindResult (readBytes 8 afterFlags)
                        (dirOffset afterDirOffset :
                          ok
                            (pair majorVersion
                              (pair minorVersion
                                (pair sectionCount
                                  (pair flags dirOffset))))
                            afterDirOffset)))))))

readSectionRecord = (bs :
  bindResult (readBytes 4 bs)
    (sectionId afterSectionId :
      bindResult (readBytes 2 afterSectionId)
        (sectionVersion afterSectionVersion :
          bindResult (readBytes 2 afterSectionVersion)
            (sectionFlags afterSectionFlags :
              bindResult (readBytes 2 afterSectionFlags)
                (compression afterCompression :
                  bindResult (readBytes 2 afterCompression)
                    (digestAlgorithm afterDigestAlgorithm :
                      bindResult (readBytes 8 afterDigestAlgorithm)
                        (offset afterOffset :
                          bindResult (readBytes 8 afterOffset)
                            (length afterLength :
                              bindResult (readBytes 32 afterLength)
                                (digest afterDigest :
                                  ok
                                    (pair sectionId
                                      (pair sectionVersion
                                        (pair sectionFlags
                                          (pair compression
                                            (pair digestAlgorithm
                                              (pair offset
                                                (pair length digest)))))))
                                    afterDigest)))))))))

readSectionDirectory_ = y (self bs sectionCount i acc :
  matchBool
    (ok (reverse acc) bs)
    (bindResult (readSectionRecord bs)
      (sectionRecord afterSectionRecord :
        self afterSectionRecord sectionCount (succ i) (pair sectionRecord acc)))
    (equal? i sectionCount))

readSectionDirectory = (sectionCount bs : readSectionDirectory_ bs sectionCount 0 t)

sectionRecordId = (sectionRecord :
  matchPair
    (sectionId _ : sectionId)
    sectionRecord)

sectionRecordVersion = (sectionRecord :
  matchPair
    (_ payload :
      matchPair
        (sectionVersion _ : sectionVersion)
        payload)
    sectionRecord)

sectionRecordFlags = (sectionRecord :
  matchPair
    (_ payload :
      matchPair
        (_ payload2 :
          matchPair
            (sectionFlags _ : sectionFlags)
            payload2)
        payload)
    sectionRecord)

sectionRecordCompression = (sectionRecord :
  matchPair
    (_ payload :
      matchPair
        (_ payload2 :
          matchPair
            (_ payload3 :
              matchPair
                (compression _ : compression)
                payload3)
            payload2)
        payload)
    sectionRecord)

sectionRecordDigestAlgorithm = (sectionRecord :
  matchPair
    (_ payload :
      matchPair
        (_ payload2 :
          matchPair
            (_ payload3 :
              matchPair
                (_ payload4 :
                  matchPair
                    (digestAlgorithm _ : digestAlgorithm)
                    payload4)
                payload3)
            payload2)
        payload)
    sectionRecord)

sectionRecordOffset = (sectionRecord :
  matchPair
    (_ payload :
      matchPair
        (_ payload2 :
          matchPair
            (_ payload3 :
              matchPair
                (_ payload4 :
                  matchPair
                    (_ payload5 :
                      matchPair
                        (offset _ : offset)
                        payload5)
                    payload4)
                payload3)
            payload2)
        payload)
    sectionRecord)

sectionRecordLength = (sectionRecord :
  matchPair
    (_ payload :
      matchPair
        (_ payload2 :
          matchPair
            (_ payload3 :
              matchPair
                (_ payload4 :
                  matchPair
                    (_ payload5 :
                      matchPair
                        (_ payload6 :
                          matchPair
                            (length _ : length)
                            payload6)
                        payload5)
                    payload4)
                payload3)
            payload2)
        payload)
    sectionRecord)

sectionRecordDigest = (sectionRecord :
  matchPair
    (_ payload :
      matchPair
        (_ payload2 :
          matchPair
            (_ payload3 :
              matchPair
                (_ payload4 :
                  matchPair
                    (_ payload5 :
                      matchPair
                        (_ payload6 :
                          matchPair
                            (_ digest : digest)
                            payload6)
                        payload5)
                    payload4)
                payload3)
            payload2)
        payload)
    sectionRecord)

lookupSectionRecord_ = y (self directory sectionId :
  matchList
    nothing
    (sectionRecord rest :
      matchBool
        (just sectionRecord)
        (self rest sectionId)
        (bytesEq? sectionId (sectionRecordId sectionRecord)))
    directory)

lookupSectionRecord = (sectionId directory : lookupSectionRecord_ directory sectionId)

sectionDirectoryHasId?_ = y (self directory sectionId :
  matchList
    false
    (sectionRecord rest :
      or?
        (bytesEq? sectionId (sectionRecordId sectionRecord))
        (self rest sectionId))
    directory)

sectionDirectoryHasId? = (sectionId directory : sectionDirectoryHasId?_ directory sectionId)

sectionDirectoryHasDuplicateIds? = y (self directory :
  matchList
    false
    (sectionRecord rest :
      or?
        (sectionDirectoryHasId?_ rest (sectionRecordId sectionRecord))
        (self rest))
    directory)

validateSectionDirectory = (directory rest :
  matchBool
    (err errDuplicateSection rest)
    (ok directory rest)
    (sectionDirectoryHasDuplicateIds? directory))

byteSlice = (offset length bytes : bytesTake length (bytesDrop offset bytes))

natMake = (bit rest :
  matchBool
    0
    (pair bit rest)
    (and? (equal? bit 0) (equal? rest 0)))

natAdd = y (self a b :
  triage
    b
    (_ : b)
    (aBit aRest :
      triage
        a
        (_ : a)
        (bBit bRest :
          matchBool
            (natMake 0 (succ (self aRest bRest)))
            (natMake (matchBool (matchBool 0 1 bBit) (matchBool 1 0 bBit) aBit)
              (self aRest bRest))
            (and? (equal? aBit 1) (equal? bBit 1)))
        b)
    a)

natDouble = (n : matchBool 0 (pair 0 n) (equal? n 0))

natTimes256 = (n :
  natDouble
    (natDouble
      (natDouble
        (natDouble
          (natDouble
            (natDouble
              (natDouble
                (natDouble n))))))))

byteNatShiftAppend_ = y (self byte acc i :
  matchBool
    acc
    (triage
      (natMake 0 (self 0 acc (succ i)))
      (_ : acc)
      (bit rest : natMake bit (self rest acc (succ i)))
      byte)
    (equal? i 8))

byteNatShiftAppend = (byte acc : byteNatShiftAppend_ byte acc 0)

beBytesToNat = (bytes :
  foldl
    (acc byte : byteNatShiftAppend byte acc)
    0
    bytes)

u32BEBytesToNat = beBytesToNat
u64BEBytesToNat = beBytesToNat

arborixHeaderMajorVersion = (header :
  matchPair
    (majorVersion _ : majorVersion)
    header)

arborixHeaderMinorVersion = (header :
  matchPair
    (_ payload :
      matchPair
        (minorVersion _ : minorVersion)
        payload)
    header)

arborixHeaderSectionCount = (header :
  matchPair
    (_ payload :
      matchPair
        (_ payload2 :
          matchPair
            (sectionCount _ : sectionCount)
            payload2)
        payload)
    header)

arborixHeaderFlags = (header :
  matchPair
    (_ payload :
      matchPair
        (_ payload2 :
          matchPair
            (_ payload3 :
              matchPair
                (flags _ : flags)
                payload3)
            payload2)
        payload)
    header)

arborixHeaderDirOffset = (header :
  matchPair
    (_ payload :
      matchPair
        (_ payload2 :
          matchPair
            (_ payload3 :
              matchPair
                (_ dirOffset : dirOffset)
                payload3)
            payload2)
        payload)
    header)

validateArborixHeader = (header rest :
  matchBool
    (ok header rest)
    (err errUnsupportedVersion rest)
    (and?
      (bytesEq? arborixMajorVersion (arborixHeaderMajorVersion header))
      (bytesEq? arborixMinorVersion (arborixHeaderMinorVersion header))))

readArborixContainer = (bs :
  bindResult (readArborixHeader bs)
    (header afterHeader :
      bindResult (validateArborixHeader header afterHeader)
        (validHeader afterValidHeader :
          bindResult (readSectionDirectory
            (u32BEBytesToNat (arborixHeaderSectionCount validHeader))
            (bytesDrop (u64BEBytesToNat (arborixHeaderDirOffset validHeader)) bs))
            (directory afterDirectory :
              bindResult (validateSectionDirectory directory afterDirectory)
                (validDirectory afterValidDirectory :
                  ok (pair validHeader validDirectory) afterValidDirectory)))))

sectionRecordOffsetNat = (sectionRecord :
  u64BEBytesToNat (sectionRecordOffset sectionRecord))

sectionRecordLengthNat = (sectionRecord :
  u64BEBytesToNat (sectionRecordLength sectionRecord))

extractSectionBytes = (sectionRecord containerBytes :
  byteSlice
    (sectionRecordOffsetNat sectionRecord)
    (sectionRecordLengthNat sectionRecord)
    containerBytes)

extractSectionBytesResult = (sectionRecord containerBytes rest :
  (sectionBytes :
    matchBool
      (ok sectionBytes rest)
      (err errUnexpectedEof rest)
      (equal? (bytesLength sectionBytes) (sectionRecordLengthNat sectionRecord)))
    (extractSectionBytes sectionRecord containerBytes))

lookupSectionBytes = (sectionId directory containerBytes :
  triage
    nothing
    (sectionRecord : just (extractSectionBytes sectionRecord containerBytes))
    (_ _ : nothing)
    (lookupSectionRecord sectionId directory))

sectionBytesOrErr = (sectionId directory containerBytes rest :
  triage
    (err errMissingSection rest)
    (sectionRecord : extractSectionBytesResult sectionRecord containerBytes rest)
    (_ _ : err errMissingSection rest)
    (lookupSectionRecord sectionId directory))

readArborixSectionBytes = (sectionId bs :
  bindResult (readArborixContainer bs)
    (container afterContainer :
      matchPair
        (_ directory : sectionBytesOrErr sectionId directory bs afterContainer)
        container))

readArborixRequiredSections = (bs :
  bindResult (readArborixContainer bs)
    (container afterContainer :
      matchPair
        (_ directory :
          bindResult (sectionBytesOrErr arborixManifestSectionId directory bs afterContainer)
            (manifestBytes _ :
              bindResult (sectionBytesOrErr arborixNodesSectionId directory bs afterContainer)
                (nodesBytes _ :
                  ok (pair manifestBytes nodesBytes) afterContainer)))
        container))

readNodeRecord = (bs :
  bindResult (readBytes 32 bs)
    (nodeHash afterNodeHash :
      bindResult (readBytes 4 afterNodeHash)
        (payloadLength afterPayloadLength :
          bindResult (readBytes (u32BEBytesToNat payloadLength) afterPayloadLength)
            (payload afterPayload :
              ok
                (pair nodeHash
                  (pair payloadLength payload))
                afterPayload))))

nodeRecordHash = (nodeRecord :
  matchPair
    (nodeHash _ : nodeHash)
    nodeRecord)

nodeRecordPayloadLength = (nodeRecord :
  matchPair
    (_ payload :
      matchPair
        (payloadLength _ : payloadLength)
        payload)
    nodeRecord)

nodeRecordPayload = (nodeRecord :
  matchPair
    (_ payload :
      matchPair
        (_ nodePayload : nodePayload)
        payload)
    nodeRecord)

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) 33))

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

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

nodePayloadStemChildHash = (nodePayload : bytesTake 32 (bytesDrop 1 nodePayload))
nodePayloadForkLeftHash = (nodePayload : bytesTake 32 (bytesDrop 1 nodePayload))
nodePayloadForkRightHash = (nodePayload : bytesTake 32 (bytesDrop 33 nodePayload))

nodeRecordPayloadValid? = (nodeRecord : nodePayloadValid? (nodeRecordPayload nodeRecord))

nodeRecordsHaveInvalidPayload? = y (self nodeRecords :
  matchList
    false
    (nodeRecord rest :
      or?
        (not? (nodeRecordPayloadValid? nodeRecord))
        (self rest))
    nodeRecords)

nodeRecordsHaveHash? = y (self nodeRecords nodeHash :
  matchList
    false
    (nodeRecord rest :
      or?
        (bytesEq? nodeHash (nodeRecordHash nodeRecord))
        (self rest nodeHash))
    nodeRecords)

nodeRecordsHaveDuplicateHashes? = y (self nodeRecords :
  matchList
    false
    (nodeRecord rest :
      or?
        (nodeRecordsHaveHash? rest (nodeRecordHash nodeRecord))
        (self rest))
    nodeRecords)

lookupNodeRecord_ = y (self nodeRecords nodeHash :
  matchList
    nothing
    (nodeRecord rest :
      matchBool
        (just nodeRecord)
        (self rest nodeHash)
        (bytesEq? nodeHash (nodeRecordHash nodeRecord)))
    nodeRecords)

lookupNodeRecord = (nodeHash nodeRecords : lookupNodeRecord_ nodeRecords nodeHash)

nodeRecordChildHashes = (nodeRecord :
  (nodePayload :
    matchBool
      t
      (matchBool
        (pair (nodePayloadStemChildHash nodePayload) t)
        (pair (nodePayloadForkLeftHash nodePayload)
          (pair (nodePayloadForkRightHash nodePayload) t))
        (nodePayloadStem? nodePayload))
      (nodePayloadLeaf? nodePayload))
    (nodeRecordPayload nodeRecord))

nodeHashPresent? = (nodeHash nodeRecords : nodeRecordsHaveHash? nodeRecords nodeHash)

nodeChildHashesPresent? = y (self childHashes nodeRecords :
  matchList
    true
    (childHash rest :
      and?
        (nodeHashPresent? childHash nodeRecords)
        (self rest nodeRecords))
    childHashes)

nodeRecordChildrenPresent? = (nodeRecord nodeRecords :
  nodeChildHashesPresent? (nodeRecordChildHashes nodeRecord) nodeRecords)

nodeRecordsClosed? = y (self nodeRecords allNodeRecords :
  matchList
    true
    (nodeRecord rest :
      and?
        (nodeRecordChildrenPresent? nodeRecord allNodeRecords)
        (self rest allNodeRecords))
    nodeRecords)

validateNodeRecords = (nodeRecords rest :
  matchBool
    (err errInvalidNodePayload rest)
    (matchBool
      (err errDuplicateNode rest)
      (matchBool
        (ok nodeRecords rest)
        (err errMissingNode rest)
        (nodeRecordsClosed? nodeRecords nodeRecords))
      (nodeRecordsHaveDuplicateHashes? nodeRecords))
    (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)))

readArborixNodesSection = (bs :
  bindResult (readArborixContainer bs)
    (container afterContainer :
      matchPair
        (_ directory :
          bindResult (sectionBytesOrErr arborixNodesSectionId directory bs afterContainer)
            (nodesBytes _ :
              bindResult (readNodesSectionComplete nodesBytes)
                (nodesSection _ : ok nodesSection afterContainer)))
        container))

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

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

nodeRecordToTreeWith = (self nodeRecords nodeRecord :
  (nodePayload :
    matchBool
      (ok t t)
      (matchBool
        (bindResult (self (nodePayloadStemChildHash nodePayload) nodeRecords)
          (child _ : ok (t child) t))
        (bindResult (self (nodePayloadForkLeftHash nodePayload) nodeRecords)
          (left _ :
            bindResult (self (nodePayloadForkRightHash nodePayload) nodeRecords)
              (right _ : ok (pair left right) t)))
        (nodePayloadStem? nodePayload))
      (nodePayloadLeaf? nodePayload))
    (nodeRecordPayload nodeRecord))

nodeHashToTree = y (self nodeHash nodeRecords :
  triage
    (err errMissingNode t)
    (nodeRecord : nodeRecordToTreeWith self nodeRecords nodeRecord)
    (_ _ : err errMissingNode t)
    (lookupNodeRecord nodeHash nodeRecords))

readArborixTreeFromHash = (rootHash bs :
  bindResult (readArborixNodesSection bs)
    (nodesSection afterContainer :
      bindResult (nodeHashToTree rootHash (nodesSectionRecords nodesSection))
        (tree _ : ok tree afterContainer)))

readArborixExecutableFromHash = readArborixTreeFromHash