Invert a Binary Tree with S assembler - https://github.com/huntiep/sasm

invertbst.sasm

(include! "syscalls.sasm")

;; Initialize malloc
(jal x1 malloc-init)

;; Create tree:
;;      a
;;     / \
;;    b   c
;;   /   / \
;;   d   e  f
(addi x10 x0 #'a')
(jal x1 create-treenode)
(add x20 x0 x10)

(addi x10 x0 #'b')
(jal x1 create-treenode)
(sd (+ x20 8) x10)
(add x21 x0 x10)

(addi x10 x0 #'d')
(jal x1 create-treenode)
(sd (+ x21 8) x10)

(addi x10 x0 #'c')
(jal x1 create-treenode)
(sd (+ x20 16) x10)
(add x21 x0 x10)

(addi x10 x0 #'e')
(jal x1 create-treenode)
(sd (+ x21 8) x10)

(addi x10 x0 #'f')
(jal x1 create-treenode)
(sd (+ x21 16) x10)


;; Print above tree in LTR order
(add x10 x0 x20)
(jal x1 print-tree)

(subi x2 x2 8)
(addi x6 x0 #'\n')
(sd x2 x6)
;; write(STDOUT, "\n", len)
(addi x10 x0 STDOUT)
(add x11 x0 x2)
(addi x12 x0 1)
(addi x17 x0 SYS_WRITE)
(ecall)


;; Invert tree
(add x10 x0 x20)
(jal x1 invert-tree)

;; Print inverted tree in LTR order
(add x10 x0 x20)
(jal x1 print-tree)

;; write(STDOUT, "\n", len)
(addi x10 x0 STDOUT)
(add x11 x0 x2)
(addi x12 x0 1)
(addi x17 x0 SYS_WRITE)
(ecall)

(addi x2 x2 8)

;; exit(0)
(add x10 x0 x0)
(addi x17 x0 SYS_EXIT)
(ecall)

;; INPUT:
;;   x10 : char
;; OUTPUT:
;;   x10 : struct TreeNode {val: char, left: *TreeNode, right: *Treenode}
create-treenode
    (subi x2 x2 16)
    (sd x2 x1)
    (sd (+ x2 8) x10)

    ;; Allocate space for treenode of 24bytes = 1(val)+7(padding)+8(left ptr)+8(right ptr)
    (addi x10 x0 24)
    (jal x1 malloc)

    (ld x1 x2)
    (ld x6 (+ x2 8))
    (addi x2 x2 16)

    ;; set tree.val
    (sd x10 x6)
    (jalr x0 x1)


;; INPUT:
;;   x10 : struct TreeNode {val u64, left: *TreeNode, right: *TreeNode}
print-tree
    ;; If tree is null return (base case)
    (beq x10 x0 print-tree-return)

    ;; Save return address and this node
    (subi x2 x2 16)
    (sd x2 x1)
    (sd (+ x2 8) x10)

    ;; Load left subtree and recurse
    (ld x10 (+ x10 8))
    (jal x1 print-tree)

    ;; Print this.val
    (ld x6 (+ x2 8))

    ;; write(STDOUT, msg, len)
    (addi x10 x0 STDOUT)
    (add x11 x0 x6)
    (addi x12 x0 1)
    (addi x17 x0 SYS_WRITE)
    (ecall)

    ;; Load right subtree and recurse
    (ld x10 (+ x6 16))
    (jal x1 print-tree)

    (ld x1 x2)
    (addi x2 x2 16)

    print-tree-return
        (jalr x0 x1)


;; INPUT:
;;   x10 : struct TreeNode {val u64, left: *TreeNode, right: *TreeNode}
;; OUTPUT:
;;   x10 : struct TreeNode {val u64, left: *TreeNode, right: *TreeNode}
invert-tree
    ;; if node == null then return
    (beq x10 x0 invert-tree-return)
    ;; save return address, root, and left subtree to stack
    (subi x2 x2 24)
    (sd x2 x1)
    (sd (+ x2 8) x10)
    (ld x11 (+ x10 8))
    (sd (+ x2 16) x11)

    ;; invert right subtree
    (ld x10 (+ x10 16))
    (jal x1 invert-tree)
    ;; set root.left to inverted root.right
    (ld x11 (+ x2 8))
    (sd (+ x11 8) x10)

    ;; invert left subtree
    (ld x10 (+ x2 16))
    (jal x1 invert-tree)
    ;; set root.right to inverted root.left
    (ld x11 (+ x2 8))
    (sd (+ x11 16) x10)
    ;; return
    (add x10 x0 x11)
    (ld x1 x2)
    (addi x2 x2 24)

    invert-tree-return
        (jalr x0 x1)

(include! "malloc.sasm")

malloc.sasm

;; Hack to avoid making global ints for now.
(define HEAP-START "\0\0\0\0\0\0\0\0")
(define HEAP-END "\0\0\0\0\0\0\0\0")
malloc-init
    (add x10 x0 x0)
    (addi x17 x0 SYS_BRK)
    (ecall)
    (la x6 HEAP-START)
    (sd x6 x10)
    (la x6 HEAP-END)
    (sd x6 x10)
    (jalr x0 x1)

;; Based on simple design outlined near the beginning of https://moss.cs.iit.edu/cs351/slides/slides-malloc.pdf
;; Args
;;   x10 - size in bytes
;; Locals
;;   x6 - temp
;;   x7 - temp
;;   x28 - temp
;;   x29 - temp
(define Malloc::Alignment 7)
(define Malloc::SIZE_T_SIZE 8)
malloc
    ;; Align(size)
    (addi x10 x10 Malloc::SIZE_T_SIZE)
    (addi x10 x10 Malloc::Alignment)
    (andi x10 x10 -8)
    (sd (- x2 8) x10)
    (sd (- x2 16) x1)
    (jal x1 find-fit)

    ;; header - x10
    ;; blk_size - x6
    (ld x1 (- x2 16))
    (beq x10 x0 malloc-allocate)
    (ld x6 (- x2 8))
    (ld x7 x10)
    (bge x7 x6 malloc-done)

    (sub x7 x7 x6)
    (add x28 x10 x6)
    (sd x28 x7)
    (jal x0 malloc-done)

    malloc-allocate
        ;; brk(CURRENT_END+blk_size)
        (la x7 HEAP-END)
        (ld x10 x7)
        (ld x6 (- x2 8))
        (add x10 x10 x6)
        (addi x17 x0 SYS_BRK)
        (ecall)
        (sd x7 x10)
        (sub x10 x10 x6)

    malloc-done
        (ori x6 x6 1)
        (sd x10 x6)
        (addi x10 x10 Malloc::SIZE_T_SIZE)
        (jalr x0 x1)

    find-fit
        (la x6 HEAP-START)
        (ld x6 x6)
        (la x7 HEAP-END)
        (ld x7 x7)
        (bge x6 x7 find-fit-null)
        (subi x10 x10 1)
        find-fit-loop
            (ld x28 x6)
            (andi x29 x28 1)
            (bne x29 x0 find-fit-set-header)
            (blt x10 x28 find-fit-done)

            (add x28 x28 x6)
            (bge x7 x28 find-fit-null)
            (ld x29 x28)
            (andi x29 x29 1)
            (bne x29 x0 find-fit-set-header)
            (ld x29 x6)
            (ld x28 x28)
            (add x28 x28 x29)
            (sd x6 x28)
            (jal x0 find-fit-loop)

            find-fit-set-header
                (xori x28 x28 1)
                (add x6 x6 x28)
                (blt x6 x7 find-fit-loop)
        find-fit-null
            (add x10 x0 x0)
            (jalr x0 x1)
        find-fit-done
            (add x10 x0 x6)
            (jalr x0 x1)

;; x10 - pointer to free
;; x6 - temporary
free
    ;; SIZE_T_SIZE = 8
    (subi x10 x10 Malloc::SIZE_T_SIZE)
    (ld x6 x10)
    (subi x6 x6 1)
    (sd x10 x6)
    (jalr x0 x1)

;; x10 - old pointer
;; x11 - new size in bytes
realloc
    (ld x6 (- x10 Malloc::SIZE_T_SIZE))
    (subi x6 x6 9)
    ;; ALIGN(size + SIZE_T_SIZE)
    (addi x11 x11 Malloc::Alignment)
    (andi x11 x11 -8)

    (bge x6 x11 realloc-done)

    (add x7 x6 x10)
    (la x28 HEAP-END)
    (ld x28 x28)
    (beq x7 x28 realloc-extend)

    (subi x2 x2 16)
    (sd (+ x2 8) x10)
    (sd x2 x1)
    (add x10 x0 x11)
    (jal x1 malloc)
    (ld x1 x2)
    (ld x11 (+ x2 8))
    (addi x2 x2 16)
    (add x6 x0 x10)

    (add x29 x0 x11)
    ;; Load the header/length of the old block
    (ld x7 (- x11 8))
    (subi x7 x7 1)
    realloc-copy-loop
        ;; Note we load a u64 at a time to save on iterations. This is okay because we
        ;; are 8-byte aligned.
        (ld x28 x11)
        (sd x6 x28)
        (addi x6 x6 8)
        (addi x11 x11 8)
        (subi x7 x7 8)
        (bne x7 x0 realloc-copy-loop)
    realloc-after-loop
        (sd (- x2 8) x10)
        (sd (- x2 16) x1)
        (add x10 x0 x29)
        (jal x1 free)
        (ld x10 (- x2 8))
        (ld x1 (- x2 16))
    realloc-done
        (jalr x0 x1)
    realloc-extend
        ;; brk(HEAP-END+new_size-old_size)
        (add x29 x0 x10)
        (add x10 x28 x11)
        (sub x10 x10 x6)
        (addi x17 x0 SYS_BRK)
        (ecall)
        (la x6 HEAP-END)
        (sd x6 x10)
        (add x10 x0 x29)
        (ori x11 x11 1)
        (sd (- x10 8) x11)
        (jalr x0 x1)

syscalls.sasm

(define STDIN 0)
(define STDOUT 1)
(define STDERR 2)

(define SYS_OPENAT 56)
(define SYS_CLOSE 57)
(define SYS_WRITE 64)
(define SYS_EXIT 93)
(define SYS_BRK 214)
(define SYS_MMAP 222)