## Copyright (C) 2021 Andrius Štikonas
## This file is part of stage0.
##
## stage0 is free software: you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation, either version 3 of the License, or
## (at your option) any later version.
##
## stage0 is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with stage0.  If not, see <http://www.gnu.org/licenses/>.

; Where the ELF Header is going to hit
; Simply jump to _start
; Our main function
:_start

; Register use:
; s2: input fd
; s3: output fd
; s4: toggle
; s5: hold
; s6: ip
; s7: tempword
; s8: shiftregister

    RD_A2 RS1_SP !16 LD               ; Input file name

    ; Initialize globals
    RD_S4 !-1 ADDI                    ; Toggle
    RD_S5 ADDI                        ; Hold
    RD_S6 ADDI                        ; Instruction Pointer

    ; Open input file and store FD in s2
    RD_A7 !56 ADDI                    ; sys_openat
    RD_A0 !-100 ADDI                  ; AT_FDCWD
    RD_A1 RS1_A2 MV                   ; file name
    RD_A2 ADDI                        ; read only
    ECALL                             ; syscall
    RS1_A0 @Fail BLTZ                 ; Error opening file
    RD_S2 RS1_A0 MV                   ; Save fd in for later

    ; Set default FD for output file to stdout
    RD_S3 !1 ADDI

    ; If we only have 2 arguments, don't use the third (it's not set)
    RD_T0 !2 ADDI
    RD_A0 RS1_SP LD                   ; Get number of the args
    RS1_A0 RS2_T0 @Fail BLT           ; No input file provided
    RS1_A0 RS2_T0 @after_open BEQ     ; No output file provided. Use stdout

    ; Open output file and store the FD in s3
    RD_A7 !56 ADDI                    ; sys_openat
    RD_A0 !-100 ADDI                  ; AT_FDCWD
    RD_A1 RS1_SP !24 LD               ; Output file (argument 3)
    RD_A2 !577 ADDI                   ; octal 00001101
    ; O_TRUNC   00001000
    ; O_CREAT   00000100
    ; O_WRONLY  00000001
    ; OCTAL!
    RD_A3 !448 ADDI                   ; Set read, write, execute permission on user
    ; S_IRWXU  00700
    ; OCTAL!
    ECALL                             ; syscall
    RD_S3 RS1_A0 MV                   ; Save fd in for later

:after_open
    RD_RA $First_pass JAL             ; First pass

    ; Rewind input file
    RD_A7 !62 ADDI                    ; sys_lseek
    RD_A0 RS1_S2 MV                   ; Input file descriptor
    RD_A1 MV                          ; Set offset to zero
    RD_A2 MV                          ; Set whence to zero
    ECALL                             ; syscall

    ; Initialize globals
    RD_S4 !-1 ADDI                    ; Toggle
    RD_S5 ADDI                        ; Hold
    RD_S6 ADDI                        ; Instruction Pointer
    RD_S7 ADDI                        ; tempword
    RD_S8 ADDI                        ; Shift register

    RD_RA $Second_pass JAL            ; Now do the second pass

    $Done JAL                         ; We are done

; First pass loop to determine addresses of labels
:First_pass
    RD_SP RS1_SP !-8 ADDI             ; Allocate stack
    RS1_SP RS2_RA SD                  ; protect ra

:First_pass_loop
    RD_RA $Read_byte JAL              ; Get another byte

    ; Deal with EOF
    RD_T1 !-4 ADDI
    RS1_A0 RS2_T1 @First_pass_done BEQ

    ; Check for :
    RD_T1 !0x3A ADDI
    RS1_A0 RS2_T1 @First_pass_0 BNE
    RD_RA $StoreLabel JAL             ; Store this label

:First_pass_0
    ; Check for !
    RD_T1 !0x21 ADDI
    RS1_A0 RS2_T1 @Throwaway_token BEQ

    ; Check for @
    RD_T1 !0x40 ADDI
    RS1_A0 RS2_T1 @Throwaway_token BEQ

    ; Check for $
    RD_T1 !0x24 ADDI
    RS1_A0 RS2_T1 @Throwaway_token BEQ

    ; Check for ~
    RD_T1 !0x7E ADDI
    RS1_A0 RS2_T1 @Throwaway_token BEQ

    RD_A1 !-1 ADDI                    ; write = false
    RD_RA $DoByte JAL                 ; Deal with everything else

    RD_T1 !-4 ADDI                    ; Deal with EOF
    RS1_A0 RS2_T1 @First_pass_done BEQ

    $First_pass_loop JAL              ; Keep looping

:Throwaway_token
    ; Deal with Pointer to label
    RD_RA $Read_byte JAL              ; Drop the char
    $First_pass_loop JAL              ; Loop again

:First_pass_done
    RD_RA RS1_SP LD                   ; restore ra
    RD_SP RS1_SP !8 ADDI              ; deallocate stack
    RS1_RA JALR                       ; return

:Second_pass
    RD_SP RS1_SP !-8 ADDI             ; Allocate stack
    RS1_SP RS2_RA SD                  ; protect ra

:Second_pass_loop
    RD_RA $Read_byte JAL              ; Read another byte

    ; Deal with EOF
    RD_T1 !-4 ADDI                    ; Deal with EOF
    RS1_A0 RS2_T1 @Second_pass_done BEQ

    ; Drop the label
    RD_T1 !0x3A ADDI
    RS1_A0 RS2_T1 @Second_pass_0 BNE

    RD_RA $Read_byte JAL              ; Read the label
    $Second_pass_loop JAL             ; Continue looping

:Second_pass_0
    ; Check for !
    RD_T1 !0x21 ADDI
    RS1_A0 RS2_T1 @UpdateShiftRegister BEQ

    ; Check for @
    RD_T1 !0x40 ADDI
    RS1_A0 RS2_T1 @UpdateShiftRegister BEQ

    ; Check for $
    RD_T1 !0x24 ADDI
    RS1_A0 RS2_T1 @UpdateShiftRegister BEQ

    ; Check for ~
    RD_T1 !0x7E ADDI
    RS1_A0 RS2_T1 @UpdateShiftRegister BEQ

    ; Deal with everything else
    RD_A1 MV                          ; write = true
    RD_RA $DoByte JAL                 ; Process our char

    # Deal with EOF
    RD_T1 !-4 ADDI
    RS1_A0 RS2_T1 @Second_pass_done BEQ ; We are done

    $Second_pass_loop JAL             ; continue looping

:UpdateShiftRegister
    RD_A1 RS1_A0 MV                   ; Store label prefix
    RD_RA $Get_table_target JAL       ; Get target
    RD_A0 RS1_A0 LD                   ; Dereference pointer
    RD_A0 RS1_A0 RS2_S6 SUB           ; target - ip

    ; Check for !
    RD_T1 !0x21 ADDI
    RS1_A1 RS2_T1 @UpdateShiftRegister_I BEQ

    ; Check for @
    RD_T1 !0x40 ADDI
    RS1_A1 RS2_T1 @UpdateShiftRegister_B BEQ

    ; Check for $
    RD_T1 !0x24 ADDI
    RS1_A1 RS2_T1 @UpdateShiftRegister_J BEQ

    ; Check for ~
    RD_T1 !0x7E ADDI
    RS1_A1 RS2_T1 @UpdateShiftRegister_U BEQ

    $Second_pass_loop JAL             ; continue looping

:UpdateShiftRegister_I
    ; Corresponds to RISC-V I format
    RD_A0 RS1_A0 !4 ADDI              ; add 4 due to this being 2nd part of AUIPC combo

    RD_T1 ~0xFFF LUI                  ; load higher bits
    RD_T1 RS1_T1 !0xFFF ADDIW
    RD_T1 RS1_A0 RS2_T1 AND           ; (value & 0xfff)
    RD_S7 RS1_T1 RS2_X20 SLLI         ; tempword = (value & 0xfff) << 20
    RD_S8 RS1_S8 RS2_S7 XOR           ; shiftregister = shiftregister ^ tempword

    $Second_pass_loop JAL             ; continue looping

:UpdateShiftRegister_B
    ; Corresponds to RISC-V B format

    ; tempword = ((value & 0x1e) << 7)            ; imm[4:1]
    ;          | ((value & 0x7e0) << (31 - 11))   ; imm[10:5]
    ;          | ((value & 0x800) >> 4)           ; imm[11]
    ;          | ((value & 0x1000) << (31 - 12))  ; imm[12]

    RD_T1 !0x1E ADDI
    RD_T1 RS1_A0 RS2_T1 AND           ; value & 0x1e
    RD_T0 RS1_T1 RS2_X7 SLLI          ; tempword = (value & 0x1e) << 7

    RD_T1 !0x7E0 ADDI
    RD_T1 RS1_A0 RS2_T1 AND           ; value & 0x7e0
    RD_T1 RS1_T1 RS2_X20 SLLI         ; (value & 0x7e0) << (31 - 11)
    RD_T0 RS1_T0 RS2_T1 OR            ; logical or with the previous expression

    RD_T1 ~0x800 LUI                  ; load higher bits
    RD_T1 RS1_T1 !0x800 ADDIW
    RD_T1 RS1_A0 RS2_T1 AND           ; value & 0x800
    RD_T1 RS1_T1 RS2_X4 SRLI          ; (value & 0x800) >> 4
    RD_T0 RS1_T0 RS2_T1 OR            ; logical or with the previous expression

    RD_T1 ~0x1000 LUI                 ; load higher bits
    RD_T1 RS1_A0 RS2_T1 AND           ; value & 0x1000
    RD_T1 RS1_T1 RS2_X19 SLLI         ; (value & 0x1000) << (31 - 12)
    RD_S7 RS1_T0 RS2_T1 OR            ; logical or with the previous expression

    RD_S8 RS1_S8 RS2_S7 XOR           ; shiftregister = shiftregister ^ tempword

    $Second_pass_loop JAL             ; continue looping

:UpdateShiftRegister_J
    ; Corresponds to RISC-V J format

    ; tempword = ((value & 0x7fe) << (30 - 10))    ; imm[10:1]
    ;          | ((value & 0x800) << (20 - 11))    ; imm[11]
    ;          | ((value & 0xff000))               ; imm[19:12]
    ;          | ((value & 0x100000) << (31 - 20)) ; imm[20]

    RD_T1 !0x7FE ADDI
    RD_T1 RS1_A0 RS2_T1 AND           ; value & 0x7fe
    RD_T0 RS1_T1 RS2_X20 SLLI         ; tempword = (value & 0x7fe) << 20

    RD_T1 ~0x800 LUI                  ; load higher bits
    RD_T1 RS1_T1 !0x800 ADDIW
    RD_T1 RS1_A0 RS2_T1 AND           ; value & 0x800
    RD_T1 RS1_T1 RS2_X9 SLLI          ; (value & 0x800) << (20 - 11)
    RD_T0 RS1_T0 RS2_T1 OR            ; logical or with the previous expression

    RD_T1 ~0xFF000 LUI                ; load higher bits
    RD_T1 RS1_A0 RS2_T1 AND           ; value & 0xff000
    RD_T0 RS1_T0 RS2_T1 OR            ; logical or with the previous expression

    RD_T1 ~0x100000 LUI               ; load higher bits
    RD_T1 RS1_A0 RS2_T1 AND           ; value & 0x100000
    RD_T1 RS1_T1 RS2_X11 SLLI         ; (value & 0x100000) << (31 - 20)
    RD_S7 RS1_T0 RS2_T1 OR            ; logical or with the previous expression

    RD_S8 RS1_S8 RS2_S7 XOR           ; shiftregister = shiftregister ^ tempword

    $Second_pass_loop JAL             ; continue looping

:UpdateShiftRegister_U
    ; Corresponds to RISC-V U format
    ; if value is 0x800 or more we have to add 11-th bit (0x1000) to compensate for signed extension

    RD_T0 ~0x800 LUI                  ; load higher bits
    RD_T0 RS1_T0 !0x800 ADDIW
    RD_T1 ~0xFFF LUI                  ; load higher bits
    RD_T1 RS1_T1 !0xFFF ADDIW

    ; We are outside 31-bit that ~ can normally load
    RD_T2 ~0x100000 LUI               ; load 0xfffff000
    RD_T2 RS1_T2 !-1 ADDIW            ; load 0xfffff000
    RD_T2 RS1_T2 RS2_X12 SLLI         ; load 0xfffff000
    RD_T1 RS1_A0 RS2_T1 AND           ; value & 0xfff
    RD_S7 RS1_A0 RS2_T2 AND           ; value & 0xfffff000
    RS1_T1 RS2_T0 @UpdateShiftRegister_U_small BLT

    # Deal with sign extension: add 0x1000
    RD_T0 ~0x1000 LUI                 ; load higher bits
    RD_S7 RS1_T0 RS2_S7 ADDW          ; (value & 0xfffff000) + 0x1000

:UpdateShiftRegister_U_small
    RD_S8 RS1_S8 RS2_S7 XOR           ; shiftregister = shiftregister ^ tempword

    $Second_pass_loop JAL             ; continue looping

:Second_pass_done
    RD_RA RS1_SP LD                   ; restore ra
    RD_SP RS1_SP !8 ADDI              ; Deallocate stack
    RS1_RA JALR                       ; return


; DoByte function
; Receives:
;   character in a0
;   bool write in a1
; Does not return anything
:DoByte
    RD_SP RS1_SP !-8 ADDI             ; Allocate stack
    RS1_SP RS2_RA SD                  ; protect ra

    RD_RA $hex JAL                    ; Process hex, store it in a6

    RS1_A6 @DoByte_Done BLTZ          ; Deal with EOF and unrecognized characters

    RS1_S4 @DoByte_NotToggle BNEZ     ; Check if toggle is set

    ; toggle = true
    RS1_A1 @DoByte_1 BNEZ             ; check if we have to write

    ; write = true
    ; We calculate (hold * 16) + hex(c) ^ sr_nextb()
    ; First, calculate new shiftregister
    RD_T0 !0xFF ADDI
    RD_T0 RS1_S8 RS2_T0 AND           ; sr_nextb = shiftregister & 0xff
    RD_S8 RS1_S8 RS2_X8 SRLI          ; shiftregister >> 8

    RD_T0 RS1_T0 RS2_A6 XOR           ; hex(c) ^ sr_nextb
    RD_T1 RS1_S5 RS2_X4 SLLI          ; hold << 4
    RD_A0 RS1_T0 RS2_T1 ADD           ; (hold << 4) + hex(c) ^ sr_nextb()
    RD_RA $fputc JAL                  ; print it
    RS1_A0 @Fail BEQZ                 ; Fail if nothing was written

:DoByte_1
    RD_S6 RS1_S6 !1 ADDI              ; Increment IP
    RD_S5 MV                          ; hold = 0
    $DoByte_FlipToggle JAL            ; return

:DoByte_NotToggle
    RD_S5 RS1_A6 MV                   ; hold = hex(c)

:DoByte_FlipToggle
    RD_S4 RS1_S4 NOT                  ; Flip the toggle

:DoByte_Done
    RD_RA RS1_SP LD                   ; restore ra
    RD_SP RS1_SP !8 ADDI              ; Deallocate stack
    RS1_RA JALR                       ; return

; Convert ASCII hex characters into binary representation, e.g. 'a' -> 0xA
; Receives:
;   character in a0
; Returns:
;   a6 with character's hex value.
:hex
    RD_SP RS1_SP !-16 ADDI            ; Allocate stack
    RS1_SP RS2_RA SD                  ; protect ra
    RS1_SP RS2_A1 @8 SD               ; protect a1

    ; Deal with EOF
    RD_T1 !-4 ADDI
    RS1_A0 RS2_T1 @hex_return BEQ

    ; deal with line comments starting with #
    RD_T1 !0x23 ADDI
    RS1_A0 RS2_T1 @ascii_comment BEQ  ; a0 eq to '#'

    ; deal with line comments starting with ;
    RD_T1 !0x3B ADDI
    RS1_A0 RS2_T1 @ascii_comment BEQ  ; a0 eq to ';'

    ; deal all ascii less than 0
    RD_T1 !0x30 ADDI
    RS1_A0 RS2_T1 @ascii_other BLT

    ; deal with 0-9
    RD_T1 !0x3A ADDI
    RS1_A0 RS2_T1 @ascii_num BLT

    ; deal with all ascii less than A
    RD_T1 !0x41 ADDI
    RS1_A0 RS2_T1 @ascii_other BLT

    ; deal with A-F
    RD_T1 !0x47 ADDI
    RS1_A0 RS2_T1 @ascii_high BLT

    ; deal with all ascii less than a
    RD_T1 !0x61 ADDI
    RS1_A0 RS2_T1 @ascii_other BLT

    ; deal with a-f
    RD_T1 !0x67 ADDI
    RS1_A0 RS2_T1 @ascii_low BLT

    ; The rest that remains needs to be ignored
    $ascii_other JAL

:ascii_num
    RD_T1 !0x30 ADDI                  ; '0' -> 0
    RD_A6 RS1_A0 RS2_T1 SUB
    $hex_return JAL                   ; return
:ascii_low
    RD_T1 !0x57 ADDI                  ; 'a' -> 0xA
    RD_A6 RS1_A0 RS2_T1 SUB
    $hex_return JAL                   ; return
:ascii_high
    RD_T1 !0x37 ADDI                  ; 'A' -> 0xA
    RD_A6 RS1_A0 RS2_T1 SUB
    $hex_return JAL                   ; return
:ascii_other
    RD_A6 !-1 ADDI                    ; Return -1
    $hex_return JAL                   ; return
:ascii_comment                        ; Read the comment until newline
    RD_RA $Read_byte JAL
    RD_T1 !0xD ADDI                   ; CR
    RS1_A0 RS2_T1 @ascii_comment_cr BEQ
    RD_T1 !0xA ADDI                   ; LF
    RS1_A0 RS2_T1 @ascii_comment BNE  ; Keep reading comment
:ascii_comment_cr
    RD_A6 !-1 ADDI                    ; Return -1
:hex_return
    RD_RA RS1_SP LD                   ; restore ra
    RD_A1 RS1_SP !8 LD                ; restore a1
    RD_SP RS1_SP !16 ADDI             ; Deallocate stack
    RS1_RA JALR                       ; return

; Read byte into a0
:Read_byte
    RD_SP RS1_SP !-16 ADDI            ; Allocate stack
    RS1_SP RS2_A1 @8 SD               ; protect a1

    RD_A7 !63 ADDI                    ; sys_read
    RD_A0 RS1_S2 MV                   ; File descriptor
    RD_A1 RS1_SP MV                   ; Get stack address for buffer
    NOP                               ; no-op
    RD_A2 !1 ADDI                     ; Size of what we want to read
    ECALL                             ; syscall

    RS1_A0 @Read_byte_1 BEQZ          ; Deal with EOF
    RD_A0 RS1_A1 LB                   ; Dereference pointer

    $Read_byte_done JAL               ; return

:Read_byte_1
    RD_A0 !-4 ADDI                    ; Put EOF in a0
:Read_byte_done
    RD_A1 RS1_SP !8 LD                ; restore a1
    RD_SP RS1_SP !16 ADDI             ; Deallocate stack
    RS1_RA JALR                       ; return

; Reads a byte and calculates table address
; Returns a pointer in a0
:Get_table_target
    RD_SP RS1_SP !-8 ADDI             ; Allocate stack
    RS1_SP RS2_RA SD                  ; protect ra

    RD_RA $Read_byte JAL              ; Get single char label
    RD_A0 RS1_A0 RS2_X3 SLLI          ; Each label in table takes 8 bytes to store
    RD_T0 ~table AUIPC                ; Load address of table
    RD_T0 RS1_T0 !table ADDI          ; into register t0
    RD_A0 RS1_A0 RS2_T0 ADD           ; Calculate offset

    RD_RA RS1_SP LD                   ; restore ra
    RD_SP RS1_SP !8 ADDI              ; Deallocate stack
    RS1_RA JALR                       ; return

:StoreLabel
    RD_SP RS1_SP !-8 ADDI             ; Allocate stack
    RS1_SP RS2_RA SD                  ; protect ra

    RD_RA $Get_table_target JAL
    RS1_A0 RS2_S6 SD                  ; Store ip into table target

    RD_RA RS1_SP LD                   ; restore ra
    RD_SP RS1_SP !8 ADDI              ; Deallocate stack
    RS1_RA JALR                       ; return

; fputc function
; Receives CHAR in a0
; Writes and returns number of bytes written in a0
:fputc
    RD_SP RS1_SP !-32 ADDI            ; allocate stack
    RS1_SP RS2_A0 SD                  ; protect a0
    RS1_SP RS2_RA @8 SD               ; protect ra
    RS1_SP RS2_A1 @16 SD              ; protect a1
    RS1_SP RS2_A2 @24 SD              ; protect a2

    RD_A7 !64 ADDI                    ; sys_write
    RD_A0 RS1_S3 MV                   ; write to output
    RD_A1 RS1_SP MV                   ; Get stack address
    RD_A2 !1 ADDI                     ; write 1 character
    ECALL                             ; syscall

    RD_RA RS1_SP !8 LD                ; restore ra
    RD_A1 RS1_SP !16 LD               ; restore a1
    RD_A2 RS1_SP !24 LD               ; restore a2
    RD_SP RS1_SP !32 ADDI             ; deallocate stack
    RS1_RA JALR                       ; return

:Done
    ; Terminate program with 0 return code
    RD_A7 !93 ADDI                    ; sys_exit
    RD_A0 MV                          ; Return code 0
    ECALL                             ; exit(0)
:Fail
    ; Terminate program with 1 return code
    RD_A7 !93 ADDI                    ; sys_exit
    RD_A0 !1 ADDI                     ; Return code 1
    ECALL                             ; exit(1)
# PROGRAM END

:table
    %0

:ELF_end
