kraken/wasm.kp

(with_import "./collections.kp"
(let (

    ; Vectors and Values
    ; Bytes encode themselves
    encode_u_LEB128 (rec-lambda recurse (x)
        (cond (< x 0x80) [x]
              true       (cons (| (& x 0x7F) 0x80) (recurse (>> x 7))))
    )
    encode_s8_LEB128  (lambda (x) (encode_u_LEB128 (& x 0xFF)))
    encode_s32_LEB128 (lambda (x) (encode_u_LEB128 (& x 0xFFFFFFFF)))
    encode_vector (lambda (enc v)
        (concat (encode_u_LEB128 (len v)) (flat_map enc v) )
    )
    encode_floating_point (lambda (x) (error "unimplemented"))
    encode_name (lambda (name)
        (encode_vector (lambda (x) [x]) name)
    )

    ; Types
    ; TODO
    encode_limits (lambda (x)
        (cond (= 1 (len x)) (concat [0x00] (encode_u_LEB128 (idx x 0)))
              (= 2 (len x)) (concat [0x01] (encode_u_LEB128 (idx x 0)) (encode_u_LEB128 (idx x 1)))
              true          (error "trying to encode bad limits"))
    )
    encode_number_type (lambda (x)
        (cond   (= x 'i32) [0x7F]
                (= x 'i64) [0x7E]
                (= x 'f32) [0x7D]
                (= x 'f64) [0x7C]
                true       (error "bad number type"))
    )
    encode_valtype (lambda (x)
        ; we don't handle reference types yet
        (encode_number_type x)
    )
    encode_result_type (lambda (x)
        (encode_vector encode_valtype x)
    )
    encode_function_type (lambda (x)
        (concat [0x60] (encode_result_type (idx x 0))
                       (encode_result_type (idx x 1)))
    )

    ; Modules
    encode_type_section (lambda (x)
        (let (
            encoded (encode_vector encode_function_type x)
        ) (concat [0x01] (encode_u_LEB128 (len encoded)) encoded ))
    )
    encode_memory_section (lambda (x)
        (let (
            encoded (encode_vector encode_limits x)
        ) (concat [0x05] (encode_u_LEB128 (len encoded)) encoded ))
    )
    encode_export (lambda (export)
        (let (
            (name type idx) export
        ) (concat (encode_name name)
                  (cond (= type 'func)   [0x00]
                        (= type 'table)  [0x01]
                        (= type 'mem)    [0x02]
                        (= type 'global) [0x03]
                        true             (error "bad export type"))
                  (encode_u_LEB128 idx)
        ))
    )
    encode_export_section (lambda (x)
        (let (
            encoded (encode_vector encode_export x)
        ) (concat [0x07] (encode_u_LEB128 (len encoded)) encoded ))
    )

    encode_function_section (lambda (x)
        (let (
            encoded (encode_vector encode_u_LEB128 x)
        ) (concat [0x03] (encode_u_LEB128 (len encoded)) encoded ))
    )
    encode_ins (lambda (ins)
        (let (
            op (idx ins 0)
        ) (cond (= op 'unreachable)            [0x00]
                (= op 'nop)                    [0x01]
                ; skipping a bunch
                ; Parametric Instructions
                (= op 'drop)                   [0x1A]
                ; skip
                ; Variable Instructions
                (= op 'local.get)      (concat [0x20] (encode_u_LEB128 (idx ins 1)))
                (= op 'local.set)      (concat [0x21] (encode_u_LEB128 (idx ins 1)))
                (= op 'local.tee)      (concat [0x22] (encode_u_LEB128 (idx ins 1)))
                (= op 'global.get)     (concat [0x23] (encode_u_LEB128 (idx ins 1)))
                (= op 'global.set)     (concat [0x24] (encode_u_LEB128 (idx ins 1)))
                ; skip
                ; Numeric Instructions
                (= op 'i32.const)      (concat [0x41] (encode_s32_LEB128 (idx ins 1)))
                ; skip
                (= op 'i32.add)                [0x6A]
        ))
    )
    encode_code (lambda (x)
        (let (
            (locals body) x
            enc_locals (encode_vector (lambda (loc)
                                    (concat (encode_u_LEB128 (idx loc 0)) (encode_valtype (idx loc 1)))) locals)
            enc_expr (concat (flat_map encode_ins body) [0x0B])
            code_bytes (concat enc_locals enc_expr)
        ) (concat (encode_u_LEB128 (len code_bytes)) code_bytes))
    )
    encode_code_section (lambda (x)
        (let (
            encoded (encode_vector encode_code x)
        ) (concat [0x0A] (encode_u_LEB128 (len encoded)) encoded ))
    )


    wasm_to_binary (lambda (wasm_code)
        (let (
            (type_section function_section memory_section export_section code_section) wasm_code
            _ (println "type_section" type_section "function_section" function_section "memory_section" memory_section "export_section" export_section "code_section" code_section)
            magic    [ 0x00 0x61 0x73 0x6D ]
            version  [ 0x01 0x00 0x00 0x00 ]
            type     (encode_type_section type_section)
            function (encode_function_section function_section)
            memory   (encode_memory_section memory_section)
            export   (encode_export_section export_section)
            code     (encode_code_section code_section)
        ) (concat magic version type function memory export code))
    )

    module (lambda (& args) (let (
        helper (rec-lambda recurse (entries i name_dict type function memory export code)
            (if (= i (len entries)) [ type function memory export code ]
                (let (
                    (n_d t f m e c) ((idx entries i) name_dict type function memory export code)
                 ) (recurse entries (+ i 1) n_d t f m e c))))
    ) (helper args 0 empty_dict [] [] [] [] [])))

    func (vau de (name & inside) (lambda (name_dict type function memory export code)
        (let (
            (params result locals body) ((rec-lambda recurse (i pe re)
                                            (cond (and (= nil pe) (< i (len inside)) (array? (idx inside i)) (< 0 (len (idx inside i))) (= 'param (idx (idx inside i) 0)))
                                                             (recurse (+ i 1) pe re)
                                                  (and (= nil pe) (= nil re) (< i (len inside)) (array? (idx inside i)) (< 0 (len (idx inside i))) (= 'result (idx (idx inside i) 0)))
                                                            ; only one result possible
                                                             (recurse (+ i 1) i (+ i 1))
                                                  (= nil pe) (recurse (+ i 1) i i)
                                                  (and            (< i (len inside)) (array? (idx inside i)) (< 0 (len (idx inside i))) (= 'local (idx (idx inside i) 0)))
                                                             (recurse (+ i 1) pe re)
                                                  true      [ (slice inside 0 (or pe 0)) (slice inside (or pe 0) (or re pe 0)) (slice inside (or re pe 0) i) (slice inside i -1) ]
                                            )
                                        ) 0 nil nil)
            result (if (!= 0 (len result)) (idx result 0)
                                           result)
            _ (println "params " params " result " result " locals " locals " body " body)
            outer_name_dict (put name_dict name (len function))
            (num_params inner_name_dict) (foldl (lambda (a x) [(+ (idx a 0) 1) (put (idx a 1) (idx x 1) (idx a 0))]) [ 0 outer_name_dict ] params)
            (num_locals inner_name_dict) (foldl (lambda (a x) [(+ (idx a 0) 1) (put (idx a 1) (idx x 1) (idx a 0))]) [ num_params inner_name_dict ] locals)
            _ (println "inner name dict" inner_name_dict)
            compressed_locals ((rec-lambda recurse (cur_list cur_typ cur_num i)
                                (cond (and (= i (len locals)) (= 0 cur_num)) cur_list
                                           (= i (len locals))                (concat cur_list [ [cur_num cur_typ] ])
                                      (= cur_typ (idx (idx locals i) 2))     (recurse cur_list                              cur_typ                (+ 1 cur_num) (+ 1 i))
                                      (= nil cur_typ)                        (recurse cur_list                              (idx (idx locals i) 2) 1             (+ 1 i))
                                      true                                   (recurse (concat cur_list [[cur_num cur_typ]]) (idx (idx locals i) 2) 1             (+ 1 i)))
                               ) [] nil 0 0)
            inner_env (add-dict-to-env de inner_name_dict)
            our_type [ (map (lambda (x) (idx x 2)) params) (slice result 1 -1) ]
            our_code (flat_map (lambda (x) (let (ins (eval x inner_env))
                                             (cond (array? ins) ins
                                                    true       (ins) ; un-evaled function, bare WAT
                                             )))
                                body)
        ) [
            outer_name_dict
            ; type
            (concat type [ our_type ])
            ; function
            (concat function [ (len function) ])
            ; memory
            memory
            ; export
            export
            ; code
            (concat code [ [ compressed_locals our_code ] ])
        ])
    ))
    i32.const (lambda (const)                                                   [['i32.const const]])
    local.get (lambda (const)                                                   [['local.get const]])
    i32.add   (lambda (& flatten)  (concat (map (lambda (x) (idx x 0)) flatten) [['i32.add]]))
    export (vau de (name t_v) (lambda (name_dict type function memory export code)
        [ name_dict type function memory (concat export [ [ name (idx t_v 0) (get-value name_dict (idx t_v 1)) ] ]) code ]
    ))
)
    (provide wasm_to_binary module func i32.const local.get i32.add export)
))