UmbyNBT

; Named Binary Tag Reader
; By Umby24
; Written for the Hypercube Minecraft Server
; ##########################################

XIncludeFile "ZLib.pbi"

#NBTBufferSize = 67108864
;{ Enumerations
Enumeration
    #NBTCompressionNone
    #NBTCompressionDetect
    #NBTCompressionGZip
    #NBTCompressionZlib
EndEnumeration

Enumeration
    #TAG_End
    #TAG_Byte
    #TAG_Short
    #TAG_Int
    #TAG_Long
    #TAG_Float
    #TAG_Double
    #TAG_Byte_Array
    #TAG_String
    #TAG_List
    #TAG_Compound
    #TAG_Int_Array
EndEnumeration
;}
;{ Structures
Structure NBTTag
    Name.s
    Type.a

    RawSize.l
    *RawData
    *Parent.NBTTag

    Byte.b
    Short.w
    Int.l
    Long.q
    Float.f
    Double.d
    String.s

    ListType.a
    ListLength.l

    List Children.NBTTag()
EndStructure
;}
;{ Helping Assembly Functions
Procedure.l Endian(val.l)
    !MOV Eax,dword[p.v_val]
    !BSWAP Eax
    !MOV dword[p.v_val], Eax
    ProcedureReturn val
EndProcedure

Procedure.f EndianF(val.f)
    !MOV Eax,dword[p.v_val]
    !BSWAP Eax
    !MOV dword[p.v_val], Eax
    ProcedureReturn val
EndProcedure

Procedure.w EndianW(val.w) ; Change Endianness of a Short (Word). Yay inline ASM!
    !MOV ax, word[p.v_val]
    !XCHG al, ah                ; Swap Lo byte <-> Hi byte
    !MOV word[p.v_val], ax
    ProcedureReturn val
EndProcedure

Procedure.q EndianQ(val.q)
    Protected Upper.l, Lower.l
    ; Split the long into two 4 byte chunks (dwords)
    PokeL(@Upper, PeekL(@val)) ; -- Upper = 1234
    PokeL(@lower, PeekL(@Val + 4)) ; Lower = 5678

    !Mov Eax, dword[p.v_Upper] ; Moves one half into Eax  -- Eax = 1234
    !Mov Ebx, dword[p.v_Lower] ; Moves the other half into Ebx -- Ebx = 5678
    !BSwap Eax ; Swap the bytes in each ; = Eax = 4321
    !BSwap Ebx ; - Ebx = 8765
    !Mov dword[p.v_Upper], Eax ; Reassign the swapped values to their variables ; - Upper = 4321
    !mov dword[p.v_Lower], Ebx ; - Lower = 8765

    PokeL(@Val, Lower) ; Push the new values, in reverse order, into the variable. ; - Val = 8765????
    PokeL(@Val + 4, Upper) ; - Val = 87654321. Swap complete.
    ProcedureReturn val
EndProcedure

Procedure.d EndianD(val.d)
    Protected Upper.l, Lower.l
    ; Split the long into two 4 byte chunks (dwords)
    PokeL(@Upper, PeekL(@val)) ; -- Upper = 1234
    PokeL(@lower, PeekL(@Val + 4)) ; Lower = 5678

    !Mov Eax, dword[p.v_Upper] ; Moves one half into Eax  -- Eax = 1234
    !Mov Ebx, dword[p.v_Lower] ; Moves the other half into Ebx -- Ebx = 5678
    !BSwap Eax ; Swap the bytes in each ; = Eax = 4321
    !BSwap Ebx ; - Ebx = 8765
    !Mov dword[p.v_Upper], Eax ; Reassign the swapped values to their variables ; - Upper = 4321
    !mov dword[p.v_Lower], Ebx ; - Lower = 8765

    PokeL(@Val, Lower) ; Push the new values, in reverse order, into the variable. ; - Val = 8765????
    PokeL(@Val + 4, Upper) ; - Val = 87654321. Swap complete.
    ProcedureReturn val
EndProcedure
;}
Global NbtTagHolder.NBTTag
Declare NbtReadCompound(*Data, *DataOffset, *Tag.NBTTag)
Declare NbtReadIntArray(*Data, *DataOffset, *Tag.NBTTag)

;{ Reading Functions
Procedure.b NbtReadByte(*Data, *DataOffset)
    Protected Offset.l, Value.b

    Offset = PeekL(*DataOffset) ; Get the value of the offset..
    Value = PeekB(*Data + Offset) ; Read the tag value..

    Offset + 1 ; Increase the offset by the size of the value
    PokeL(*DataOffset, Offset) ; Push the new offset size to memory

    ProcedureReturn Value ; And return the value.
EndProcedure

Procedure.w NbtReadShort(*Data, *DataOffset)
    Protected Offset.l, Value.w

    Offset = PeekL(*DataOffset) ; Get the value of the offset..
    Value = EndianW(PeekW(*Data + Offset)) ; Read the tag value..

    Offset + 2 ; Increase the offset by the size of the value
    PokeL(*DataOffset, Offset) ; Push the new offset size to memory

    ProcedureReturn Value ; And return the value.
EndProcedure

Procedure.l NbtReadInt(*Data, *DataOffset)
    Protected Offset.l, Value.l

    Offset = PeekL(*DataOffset) ; Get the value of the offset..
    Value = Endian(PeekL(*Data + Offset)) ; Read the tag value..

    Offset + 4 ; Increase the offset by the size of the value
    PokeL(*DataOffset, Offset) ; Push the new offset size to memory

    ProcedureReturn Value ; And return the value.
EndProcedure

Procedure.q NbtReadLong(*Data, *DataOffset)
    Protected Offset.l, Value.q

    Offset = PeekL(*DataOffset) ; Get the value of the offset..
    Value = EndianQ(PeekQ(*Data + Offset)) ; Read the tag value..

    Offset + 8 ; Increase the offset by the size of the value
    PokeL(*DataOffset, Offset) ; Push the new offset size to memory

    ProcedureReturn Value ; And return the value.
EndProcedure

Procedure.f NbtReadFloat(*Data, *DataOffset)
    Protected Offset.l, Value.f
    Offset = PeekL(*DataOffset) ; Get the value of the offset..
    Value = EndianF(PeekF(*Data + Offset)) ; Read the tag value..

    Offset + 4 ; Increase the offset by the size of the value
    PokeL(*DataOffset, Offset) ; Push the new offset size to memory

    ProcedureReturn Value ; And return the value.
EndProcedure

Procedure.d NbtReadDouble(*Data, *DataOffset)
    Protected Offset.l, Value.d

    Offset = PeekL(*DataOffset) ; Get the value of the offset..
    Value = EndianD(PeekD(*Data + Offset)) ; Read the tag value..

    Offset + 8 ; Increase the offset by the size of the value
    PokeL(*DataOffset, Offset) ; Push the new offset size to memory

    ProcedureReturn Value ; And return the value.
EndProcedure

Procedure NbtReadByteArray(*Data, *DataOffset, *Tag.NBTTag)
    Protected Offset.l, Size.l

    Size = NbtReadInt(*Data, *DataOffset)
    Offset = PeekL(*DataOffset) ; Get the value of the offset..

    *Tag\RawData = AllocateMemory(Size)
    *Tag\RawSize = Size
    CopyMemory(*Data + Offset, *Tag\RawData, Size)

    Offset + Size
    PokeL(*DataOffset, Offset) ; Push the new offset size to memory
EndProcedure

Procedure.s NbtReadString(*Data, *DataOffset)
    Protected Offset.l, Size.l, Value.s
    Offset = PeekL(*DataOffset)

    Size = EndianW(PeekW(*Data + Offset))
    Offset + 2
    Value = PeekS(*Data + Offset, Size)
    Offset + Size
    PokeL(*DataOffset, Offset)

    ProcedureReturn Value
EndProcedure

Procedure NbtReadList(*Data, *DataOffset, *Tag.NBTTag)
    Protected Offset.l, Size.l, Type.l, i.l
    Offset = PeekL(*DataOffset)
    Type = PeekA(*Data + Offset)
    Offset + 1
    PokeL(*DataOffset, Offset)

    ; - Get list size.
    Size = NbtReadInt(*Data, *DataOffset)
    *Tag\ListType = Type
    *Tag\ListLength = Size

    For i = 1 To Size ; Iterate the list.
        Define *ChildElement.NBTTag = AddElement(*Tag\Children())
        *Tag\Children()\Type = Type
        *Tag\Children()\Name = "[None]"

        Select Type
            Case #TAG_Byte
                *Tag\Children()\Byte = NbtReadByte(*Data, *DataOffset)
            Case #TAG_Short
                *Tag\Children()\Short = NbtReadShort(*Data, *DataOffset)
            Case #TAG_Int
                *Tag\Children()\Int = NbtReadInt(*Data, *DataOffset)
            Case #TAG_Long
                *Tag\Children()\Long = NbtReadLong(*Data, *DataOffset)
            Case #TAG_Float
                *Tag\Children()\Float = NbtReadFloat(*Data, *DataOffset)
            Case #TAG_Double
                *Tag\Children()\Double = NbtReadDouble(*Data, *DataOffset)
            Case #TAG_Byte_Array
                NbtReadByteArray(*Data, *DataOffset, *Tag\Children())
            Case #TAG_String
                *Tag\Children()\String = NbtReadString(*Data, *DataOffset)
            Case #TAG_List
                NbtReadList(*Data, *DataOffset, *ChildElement)
            Case #TAG_Compound
                NbtReadCompound(*Data, *DataOffset, *ChildElement)
            Case #TAG_Int_Array
                NbtReadIntArray(*Data, *DataOffset, *ChildElement)
        EndSelect
    Next

EndProcedure

Procedure NbtReadIntArray(*Data, *DataOffset, *Tag.NBTTag)
    Protected Offset.l, Size.l, i.l, RawOffset.l
    Size = NbtReadInt(*Data, *DataOffset)
    ;Offset = PeekL(*DataOffset)

    *Tag\RawData = AllocateMemory(Size * 4)
    *Tag\RawSize = Size

    For i = 1 To Size
        ;PokeL(*DataOffset, Offset)
        PokeL(*Tag\RawData + RawOffset, NbtReadInt(*Data, *DataOffset))
        RawOffset + 4
    Next

EndProcedure

Procedure NbtReadCompound(*Data, *DataOffset, *Tag.NBTTag)
    Protected Offset.l
    Offset = PeekL(*DataOffset)

    While 1=1
        TagType.a = PeekA(*Data + Offset)
        Offset+1

        If TagType <> #TAG_End
            PokeL(*DataOffset, Offset)
            AddElement(*Tag\Children())
            *Tag\Children()\Type = TagType
            *Tag\Children()\Name = NbtReadString(*Data, *DataOffset)
            *Tag\Children()\Parent = *Tag
        EndIf

        Select TagType
            Case #TAG_End
                Break ; -- Break out of the reading loop.
            Case #TAG_Byte
                *Tag\Children()\Byte = NbtReadByte(*Data, *DataOffset)
            Case #TAG_Short
                *Tag\Children()\Short = NbtReadShort(*Data, *DataOffset)
            Case #TAG_Int
                *Tag\Children()\Int = NbtReadInt(*Data, *DataOffset)
            Case #TAG_Long
                *Tag\Children()\Long = NbtReadLong(*Data, *DataOffset)
            Case #TAG_Float
                *Tag\Children()\Float = NbtReadFloat(*Data, *DataOffset)
            Case #TAG_Double
                *Tag\Children()\Double = NbtReadDouble(*Data, *DataOffset)
            Case #TAG_Byte_Array
                NbtReadByteArray(*Data, *DataOffset, *Tag\Children())
            Case #TAG_String
                *Tag\Children()\String = NbtReadString(*Data, *DataOffset)
            Case #TAG_List
                NbtReadList(*Data, *DataOffset, *Tag\Children())
            Case #TAG_Compound
                NbtReadCompound(*Data, *DataOffset, *Tag\Children())
            Case #TAG_Int_Array
                NbtReadIntArray(*Data, *DataOffset, *Tag\Children())
        EndSelect

        Offset = PeekL(*DataOffset)
    Wend

    PokeL(*DataOffset, Offset)
EndProcedure

Procedure ParseNbtData(*YourTag.NBTTag, *Data, Datalength)
    Define Offset.l

    TagType.a = PeekA(*Data + Offset)
    Offset+1

    If TagType <> #TAG_Compound
        PrintN("Compound tag not the base.")
        ProcedureReturn
    EndIf

    *YourTag\Type = TagType
    *YourTag\Name = NbtReadString(*Data, @Offset)
    NbtReadCompound(*Data, @Offset, *YourTag)
EndProcedure

Procedure DecompressData(*Data, DataLength)
    Protected Stream.z_stream, BufferSize

    BufferSize = #NBTBufferSize
    *Temp = AllocateMemory(BufferSize)

    Stream\avail_in = DataLength
    Stream\avail_out = BufferSize
    Stream\next_in = *Data
    Stream\next_out = *Temp

    If Not inflateInit2_(@Stream, 47, zlibVersion(), SizeOf(z_stream)) = #Z_OK
        FreeMemory(*Temp)
        ProcedureReturn #Null
    EndIf

    Repeat
        TempResult = inflate(Stream, #Z_NO_FLUSH)

        If TempResult <> #Z_OK And TempResult <> #Z_STREAM_END
            FreeMemory(*Temp)
            inflateEnd(Stream)
            ProcedureReturn #Null
        EndIf

        If TempResult = #Z_OK
            BufferSize + #NBTBufferSize
            *Temp = ReAllocateMemory(*Temp, BufferSize)

            Stream\avail_out = #NBTBufferSize
            Stream\next_out = *Temp + BufferSize - #NBTBufferSize
        EndIf
    Until TempResult = #Z_STREAM_END

    inflateEnd(Stream)

    ProcedureReturn *Temp
EndProcedure

Procedure ParseNBT(*Data, DataLength, CompressionMode)
    ; This file will first decompress the data (if need be) then pass it off for real parsing, and then return it.
    Select CompressionMode
        Case #NBTCompressionNone
            ParseNbtData(@NbtTagHolder, *Data, DataLength)
            FreeMemory(*Data)
        Case #NBTCompressionDetect
            If PeekA(*Data) = $1F And PeekA(*Data + 1) = $8B
                ; Compressed. Need to decompress.
                *Uncompressed = DecompressData(*Data, DataLength)
                ; We can now free our compressed data from memory, it is no longer needed.
                FreeMemory(*Data)

                If *Uncompressed <> #Null
                    ParseNbtData(@NbtTagHolder, *Uncompressed, MemorySize(*Uncompressed))
                    FreeMemory(*Uncompressed)
                EndIf

            Else
                ParseNbtData(@NbtTagHolder, *Data, DataLength)
                FreeMemory(*Data)
            EndIf
        Case #NBTCompressionGZip, #NBTCompressionZlib
            *Uncompressed = DecompressData(*Data, DataLength)
            ; We can now free our compressed data from memory, it is no longer needed.
            FreeMemory(*Data)

            If *Uncompressed <> #Null
                ParseNbtData(@NbtTagHolder, *Uncompressed, MemorySize(*Uncompressed))
            EndIf

            FreeMemory(*Uncompressed)
    EndSelect

    ProcedureReturn @NbtTagHolder
EndProcedure

Procedure ReadNBTFile(Filename.s, CompressionMode=#NBTCompressionDetect)
    Protected FileID, *FileBuffer

    FileID = OpenFile(#PB_Any, Filename)

    If Not FileID
        ProcedureReturn ; Couldn't open file.
    EndIf

    *FileBuffer = AllocateMemory(Lof(FileID))

    If Not *FileBuffer
        CloseFile(FileID)
        ProcedureReturn ; Failed to allocate memory
    EndIf

    If ReadData(FileID, *FileBuffer, Lof(FileID)) <> Lof(FileID)
        CloseFile(FileID)
        FreeMemory(*FileBuffer)
        ProcedureReturn
    EndIf

    CloseFile(FileID)

    ProcedureReturn ParseNBT(*FileBuffer, MemorySize(*FileBuffer), CompressionMode)
EndProcedure
;}
;{ Saving Functions
Procedure GetNBTSize(*Tag.NBTTag)
    Protected Size.l

    If Not *Tag
        ProcedureReturn #False
    EndIf

    If Not (*Tag\Parent And *Tag\Parent\Type = #TAG_List)
        Size + 3 + StringByteLength(*Tag\Name, #PB_UTF8)
    EndIf

    ;Get payload size
    Select *Tag\Type
        Case #TAG_Byte : Size + 1
        Case #TAG_Short : Size + 2
        Case #TAG_Int : Size + 4
        Case #TAG_Long : Size + 8
        Case #TAG_Float : Size + 4
        Case #TAG_Double : Size + 8
        Case #TAG_Byte_Array : Size + 4 + *Tag\RawSize
        Case #TAG_String : Size + 2 + StringByteLength(*Tag\String, #PB_UTF8)
        Case #TAG_List
            Size + 5 ;Type and list size.

            ForEach *Tag\Children()
                Size + GetNBTSize(*Tag\Children())
            Next
        Case #TAG_Compound
            ForEach *Tag\Children()
                Size + GetNBTSize(*Tag\Children())
            Next
            Size + 1 ; Tag_End
        Case #TAG_Int_Array : Size + 4 + (*Tag\RawSize * 4)
    EndSelect

    ProcedureReturn Size
EndProcedure
;
Procedure SaveNBT(*Tag.NBTTag, *Memory)
    If Not *Tag
        ProcedureReturn #False
    EndIf

    If Not *Memory
        ProcedureReturn #False
    EndIf

    If Not (*Tag\Parent And *Tag\Parent\Type = #TAG_List)
        PokeB(*Memory, *Tag\Type) : *Memory + 1 ; Tag Type
        PokeW(*Memory, EndianW(StringByteLength(*Tag\Name, #PB_UTF8))) : *Memory + 2 ; Length of name
        PokeS(*Memory, *Tag\Name, -1, #PB_UTF8) : *Memory + StringByteLength(*Tag\Name, #PB_UTF8) ; The actual name
    EndIf

    ;Get payload size
    Select *Tag\Type
        Case #TAG_Byte
            PokeB(*Memory, *Tag\Byte)
            *Memory + 1
        Case #TAG_Short
            PokeW(*Memory, EndianW(*Tag\Short))
            *Memory + 2
        Case #TAG_Int
            PokeL(*Memory, Endian(*Tag\Int))
            *Memory + 4
        Case #TAG_Long
            PokeQ(*Memory, EndianQ(*Tag\Long))
            *Memory + 8
        Case #TAG_Float
            PokeF(*Memory, EndianF(*Tag\Float))
            *Memory + 4
        Case #TAG_Double
            PokeD(*Memory, EndianD(*Tag\Double))
            *Memory + 8
        Case #TAG_Byte_Array
            PokeL(*Memory, Endian(*Tag\RawSize))
            *Memory + 4

            If *Tag\RawData
                CopyMemory(*Tag\RawData, *Memory, *Tag\RawSize)
            EndIf

            *Memory + *Tag\RawSize
        Case #TAG_String
            PokeW(*Memory, EndianW(*Tag\RawSize))
            *Memory + 2

            PokeS(*Memory, *Tag\String, -1, #PB_UTF8)
            *Memory + *Tag\RawSize
        Case #TAG_List
            PokeB(*Memory, *Tag\ListType) : *Memory + 1
            PokeL(*Memory, Endian(*Tag\ListLength))
            *Memory + 4

            ForEach *Tag\Children()
                *Memory = SaveNBT(*Tag\Children(), *Memory)
            Next

        Case #TAG_Compound
            ForEach *Tag\Children()
                *Memory = SaveNBT(*Tag\Children(), *Memory)
            Next

            PokeB(*Memory, 0) : *Memory + 1
        Case #TAG_Int_Array
            PokeL(*Memory, Endian(*Tag\RawSize)) : *Memory + 4
            Define i.l, offset.l
            offset = 0

            For i = 1 To *Tag\RawSize
                PokeL(*Memory, Endian(PeekL(*Tag\RawData + offset)))
                *Memory + 4
                offset + 4
            Next

            *Memory + (*Tag\RawSize * 4)
    EndSelect

    ProcedureReturn *Memory
EndProcedure
;
Procedure SaveNBTFile(*NBT.NBTTag, Filename.s, CompressionMode=#NBTCompressionGZip)
    Protected File.l, *NbtMem, NbtSize.l

    If Not *NBT
        ProcedureReturn #False
    EndIf

    File = CreateFile(#PB_Any, Filename)

    If Not File
        ProcedureReturn #False
    EndIf

    NbtSize = GetNBTSize(*NBT)
    *NbtMem = AllocateMemory(NbtSize)

    PrintN(Str(NbtSize))
    PrintN(Str(*NbtMem))

    If Not *NbtMem
        CloseFile(File)
        ProcedureReturn #False
    EndIf

    SaveNBT(*NBT, *NbtMem)

    Select CompressionMode
        Case #NBTCompressionNone
            WriteData(File, *NbtMem, NbtSize)
        Case #NBTCompressionGZip
        Case #NBTCompressionZlib
    EndSelect

    CloseFile(File)
    ProcedureReturn #True
EndProcedure

;}
;{ Manipulation Functions
Procedure AddTag(*Tag.NBTTag, TagType, TagName)
;}
;{ Serialization to string
Procedure.s SerializeNbt(*Tag.NBTTag, Level=0)
    Protected Output.s = Space(Level*4)

    If Not *Tag
        ProcedureReturn "[Invalid_Element]"
    EndIf

    Select *Tag\Type
        Case #TAG_End
            Output + "[This shouldn't appear!]"
        Case #TAG_Byte
            Output + "TAG_Byte('"+*Tag\Name+"'): "+Str(*Tag\Byte)
        Case #Tag_Short
            Output + "TAG_Short('"+*Tag\Name+"'): "+Str(*Tag\Short)
        Case #TAG_Int
            Output + "TAG_Int('"+*Tag\Name+"'): "+Str(*Tag\Int)
        Case #TAG_Long
            Output + "TAG_Long('"+*Tag\Name+"'): "+Str(*Tag\Long)
        Case #TAG_Float
            Output + "TAG_Float('"+*Tag\Name+"'): "+StrF(*Tag\Float)
        Case #TAG_Double
            Output + "TAG_Double('"+*Tag\Name+"'): "+StrD(*Tag\Double)
        Case #Tag_Byte_Array
            Output + "TAG_Byte_Array('"+*Tag\Name+"'): ["+Str(*Tag\RawSize)+" bytes]"
        Case #Tag_String
            Output + "TAG_String('"+*Tag\Name+"'): '"+*Tag\String+"'"
        Case #Tag_List
            Output + "TAG_List('"+*Tag\Name+"'): "+Str(ListSize(*Tag\Children()))+" entries" + #CRLF$
            Output + Space(Level*4) + "{" + #CRLF$

            ForEach *Tag\Children()
                Output + SerializeNbt(*Tag\Children(), Level+1) + #CRLF$
            Next

            Output + Space(Level*4) + "}"
        Case #Tag_Compound
            Output + "TAG_Compound('"+*Tag\Name+"'): "+Str(ListSize(*Tag\Children()))+" entries" + #CRLF$
            Output + Space(Level*4) + "{" + #CRLF$

            ForEach *Tag\Children()
                Output + SerializeNbt(*Tag\Children(), Level+1) + #CRLF$
            Next

            Output + Space(Level*4) + "}"
        Case #Tag_Int_Array
            Output + "TAG_Int_Array('"+*Tag\Name+"'): ["+Str(*Tag\RawSize)+" ints]"
        Default
            Output + "[Something went really wrong here!]"
    EndSelect

    ProcedureReturn Output
EndProcedure

;}