;; my wasm reader
(defparameter endian :little)
(defstruct env customs types imports exports memory funcs globals codes data start)
(defparameter verbose nil)

(defun read-u32 (stream)
  (let* ((bytes (make-array 4 :element-type '(unsigned-byte 8) :initial-element 0))
         (end (read-sequence bytes stream)))
    (assert (= 4 end))
    (reduce #'(lambda (acc el) (logior (ash acc 8) el)) (reverse bytes) :initial-value 0)))

(defun read-leb (stream &key (word-bits 32))
  (iter (with acc = 0) (with nbytes = 0) (with maxnbytes = (ceiling (/ word-bits 7)))
    (for bytecnt from 0 below maxnbytes)
    (for bitpos upfrom 0 by 7)
    (for x = (read-byte stream nil nil))
    (while x)
    (incf nbytes)
    (setf acc (dpb x (byte 8 bitpos) acc))
    (while (/= 0 (logand #x80 x)))
    (finally
      (let* ((full-word-mask (1- (expt 2 word-bits))))
        (return (values (logand full-word-mask acc) nbytes))))))

(defun read-sleb (stream &key (word-bits 32))
  (iter
    (with nbytes = 0) (with byte = 0) (with result = 0) (with shift = 0)
    (with full-word-mask = (1- (expt 2 word-bits)))
    (with mid-value = (expt 2 (1- word-bits)))
    (setf byte (read-byte stream nil nil))
    (while byte)
    (incf nbytes)
    (setf result (logior result (ash (logand #x7f byte) shift)))
    (incf shift 7)
    (while (/= 0 (logand #x80 byte)))
    (finally
      (if (and (< shift word-bits) (/= 0 (logand #x40 byte)))
        (setf result (logior result (ash #xffffffff shift))))
      (setf result (logand result full-word-mask))
      (if (>= result mid-value)
        (setf result (- result full-word-mask 1)))
      (return (values result nbytes)))))

(defun expect-magic (nbytes magic stream)
  (ecase nbytes
    (1 (let* ((u8 (read-byte stream nil nil)))
         (unless (= magic u8) (error "bad magic expected:#x~2,'0x got:#x~2,'0x" magic u8))))
    (4 (let* ((u32 (read-u32 stream)))
         (unless (= magic u32) (error "bad magic expected:#x~8,'0x got:#x~8,'0x" magic u32))))))

(defun read-boolean (stream)
  (/= 0 (read-byte stream nil nil)))

(defun code-to-valtype (n)
  (aref #(f64 f32 i64 i32) (- n #x7c)))
(defun read-valtype (stream)
  (code-to-valtype (read-byte stream nil nil)))
(defun read-valtype-vec (stream)
  (iter (with nels = (read-leb stream))
    (for i below nels)
    (collect (read-valtype stream))))
(defun read-wasm-byte-vec (stream)
  (let* ((nbytes (read-leb stream))
         (bytes (make-array nbytes :element-type '(unsigned-byte 8) :initial-element 0))
         (end (read-sequence bytes stream)))
    (assert (= end nbytes))
    bytes))
(defun read-wasm-utf8-string (stream)
  (babel:octets-to-string (read-wasm-byte-vec stream) :encoding :utf-8))
(defun read-result-type (stream) ;; todo what's different about this signed version s33 v u32 read-leb
  (let* ((result-type (read-sleb stream)))
    (case result-type
      (#x40 'no-result)
      ((#x7c #x7d #x7e #x7f) (code-to-valtype result-type))
      (otherwise result-type))))

(defparameter section-tag-names #(
  custom-section ;; 0
  type-section ;; 1
  import-section ;; 2
  function-section ;; 3
  table-section ;; 4
  memory-section ;; 5
  global-section ;; 6
  export-section ;; 7
  start-section ;; 8
  element-section ;; 9
  code-section ;; 10
  data-section ;; 11
  ))

(defun section-tag-from-byte (u8) (aref section-tag-names u8))

(defun read-custom-section (size stream env)
  (let* ((blob (make-array size :element-type '(unsigned-byte 8) :initial-element 0))
         (end (read-sequence blob stream)))
    (assert (= end size))
    (push blob (env-customs env))))
(defun read-type-section (size stream env)
  (declare (ignore size))
  (if (env-types env) (error "type-section only allowed once, found a second one"))
  (let* ((nels (read-leb stream)))
    (setf (env-types env)
      (iter (for i below nels)
        (expect-magic 1 #x60 stream)
        (collect
          (list :params (read-valtype-vec stream) :results (read-valtype-vec stream)))))))

(defun read-limits (stream)
  (let* ((min-max (read-boolean stream)))
    `(:limits ,(read-leb stream) ,@(if min-max (list (read-leb stream))))))
(defun read-import-desc (stream)
  (let* ((tag (aref #(:function :table :memory :global) (read-byte stream nil nil))))
    (list
      tag
      (ecase tag
        (:function (read-leb stream))
        (:memory (read-limits stream))
        ;; todo
        ))))
(defun read-import-section (size stream env)
  (declare (ignore size))
  (if (env-imports env) (error "import-section only allowed once, found a second one"))
  (let* ((nels (read-leb stream)))
    (setf
      (env-imports env)
      (iter (for i below nels)
        (collect
          (list :module (read-wasm-utf8-string stream) :name (read-wasm-utf8-string stream)
                :desc (read-import-desc stream)))))))

(defun read-function-section (size stream env)
  (declare (ignore size))
  (if (env-funcs env) (error "function-section only allowed once, found a second one"))
  (let* ((nels (read-leb stream)))
    (setf (env-funcs env)
      (cons :typeidx
        (iter (for i below nels)
          (collect (read-leb stream)))))))

(defun read-memory-section (size stream env)
  (declare (ignore size))
  (if (env-memory env) (error "memory-section only allowed once, found a second one"))
  (let* ((nels (read-leb stream)))
    (setf (env-memory env)
      (iter (for i below nels)
        (collect (read-limits stream))))))

(defparameter instr-name (make-array 256 :initial-element nil))
(defparameter instr-handler (make-array 256 :initial-element nil))

(declaim (ftype (function (stream) (values list integer)) read-expr))

(defun read-instr-solo (opcode stream)
  (declare (ignore stream))
  (list opcode))
(defun read-instr-imm (opcode stream)
  (list opcode (read-leb stream)))
(defun read-instr-signed-imm (opcode stream)
  (list opcode
    (ecase opcode
      (i32.const (read-sleb stream :word-bits 32))
      (i64.const (read-sleb stream :word-bits 64)))))

(defun read-instr-2ximm (opcode stream)
  (list opcode (read-leb stream) (read-leb stream)))
(defun read-instr-zero (opcode stream)
  (unless (zerop (read-byte stream nil nil)) (error "opcode ~a expected zero byte suffix" opcode))
  (list opcode))
(defun read-block-or-loop-expr (opcode stream)
  (let* ((result-type (read-result-type stream)))
    (multiple-value-bind (expr-seq terminator) (read-expr stream)
      (unless (= #x0b terminator) (error "bad block-expr terminator. Expected #x0b but got #x~2,'0x" terminator))
      `(,opcode ,result-type ,expr-seq))))
(defun read-if-expr (opcode stream)
  (declare (ignore opcode))
  (let* ((result-type (read-result-type stream)))
    (multiple-value-bind (then-seq terminator) (read-expr stream)
      (cond
        ((= #x05 terminator)
          (multiple-value-bind (else-seq terminator) (read-expr stream)
            (unless (= #x0b terminator) (error "bad if-else-expr terminator. Expected #x0b but got #x~2,'0x" terminator))
            `(if-then-else ,result-type ,then-seq ,else-seq)))
        ((= #x0b terminator)
          `(if-then ,result-type ,then-seq))
        (t (error "bad if-then-expr terminator. Expected #x0b but got #x~2,'0x" terminator))))))
(defun init-instr-lookup ()
  (iter
    (for i upfrom #x45)
    (for instr in-vector
      #(i32.eqz i32.eq i32.ne i32.lt_s i32.lt_u i32.gt_s i32.gt_u i32.le_s i32.le_u i32.ge_s i32.ge_u i64.eqz i64.eq i64.ne i64.lt_s i64.lt_u
      i64.gt_s i64.gt_u i64.le_s i64.le_u i64.ge_s i64.ge_u f32.eq f32.ne f32.lt f32.gt f32.le f32.ge f64.eq f64.ne f64.lt f64.gt f64.le
      f64.ge i32.clz i32.ctz i32.popcnt i32.add i32.sub i32.mul i32.div_s i32.div_u i32.rem_s i32.rem_u i32.and i32.or i32.xor i32.shl
      i32.shr_s i32.shr_u i32.rotl i32.rotr i64.clz i64.ctz i64.popcnt i64.add i64.sub i64.mul i64.div_s i64.div_u i64.rem_s i64.rem_u
      i64.and i64.or i64.xor i64.shl i64.shr_s i64.shr_u i64.rotl i64.rotr f32.abs f32.neg f32.ceil f32.floor f32.trunc f32.nearest
      f32.sqrt f32.add f32.sub f32.mul f32.div f32.min f32.max f32.copysign f64.abs f64.neg f64.ceil f64.floor f64.trunc f64.nearest
      f64.sqrt f64.add f64.sub f64.mul f64.div f64.min f64.max f64.copysign i32.wrap_i64 i32.trunc_f32_s i32.trunc_f32_u i32.trunc_f64_s
      i32.trunc_f64_u i64.extend_i32_s i64.extend_i32_u i64.trunc_f32_s i64.trunc_f32_u i64.trunc_f64_s i64.trunc_f64_u f32.convert_i32_s
      f32.convert_i32_u f32.convert_i64_s f32.convert_i64_u f32.demote_f64 f64.convert_i32_s f64.convert_i32_u f64.convert_i64_s
      f64.convert_i64_u f64.promote_f32 i32.reinterpret_f32 i64.reinterpret_f64 f32.reinterpret_i32 f64.reinterpret_i64 i32.extend8_s
      i32.extend16_s i64.extend8_s i64.extend16_s i64.extend32_s))
    (setf (aref instr-name i) instr)
    (setf (aref instr-handler i) #'read-instr-solo))
  (iter
    (for (i instr) in '(
      (#x00 unreachable)
      (#x01 nop)
      (#x05 else)
      (#x0b end)
      (#x0f return)
      (#x1a drop)
      (#x1b select)))
    (setf (aref instr-name i) instr)
    (setf (aref instr-handler i) #'read-instr-solo))
  (iter
    (for (i instr) in '(
      (#x3f memory.size)
      (#x40 memory.grow)))
    (setf (aref instr-name i) instr)
    (setf (aref instr-handler i) #'read-instr-zero))
  (iter
    (for i upfrom #x20)
    (for instr in '(local.get local.set local.tee global.get global.set))
    (setf (aref instr-name i) instr)
    (setf (aref instr-handler i) #'read-instr-imm))
  (iter
    (for i upfrom #x41)
    (for instr in '(i32.const i64.const f32.const f64.const))
    (setf (aref instr-name i) instr)
    (setf (aref instr-handler i) #'read-instr-signed-imm))
  (iter
    (for (i instr) in '(
      (#x0c br)
      (#x0d br_if)
      (#x10 call)))
    (setf (aref instr-name i) instr)
    (setf (aref instr-handler i) #'read-instr-imm))
  (iter
    (for i upfrom #x28)
    (for instr in '(
      i32.load i64.load f32.load f64.load i32.load8_s i32.load8_u i32.load16_s i32.load16_u i64.load8_s i64.load8_u i64.load16_s
      i64.load16_u i64.load32_s i64.load32_u i32.store i64.store f32.store f64.store i32.store8 i32.store16 i64.store8 i64.store16 i64.store32))
    (setf (aref instr-name i) instr)
    (setf (aref instr-handler i) #'read-instr-2ximm))
  (setf (aref instr-name #x02) 'block (aref instr-handler #x02) #'read-block-or-loop-expr)
  (setf (aref instr-name #x03) 'loop (aref instr-handler #x03) #'read-block-or-loop-expr)
  (setf (aref instr-name #x04) 'if (aref instr-handler #x04) #'read-if-expr))

(defun read-expr (stream)
  (iter
    (for opcode = (read-byte stream nil nil))
    (while opcode)
    (if (position opcode #(#x0b #x05)) (finish))
    (if (null (aref instr-handler opcode)) (error "opcode #x~2,'0x undefined~%" opcode))
    (collect (funcall (aref instr-handler opcode) (aref instr-name opcode) stream) into expr)
    (finally (return (values expr opcode)))
    ))

(defun read-global-section (size stream env)
  (declare (ignore size))
  (if (env-globals env) (error "global-section only allowed once, found a second one"))
  (let* ((nels (read-leb stream)))
    (setf (env-globals env)
      (iter (for i below nels)
        (collect `(:type ,(read-valtype stream) :mutable ,(read-boolean stream) :value ,(car (read-expr stream))))))))

(defun read-export-tag (stream)
  (aref #(:funcidx :tableidx :memidx :globalidx) (read-byte stream nil nil)))
(defun read-export-section (size stream env)
  (declare (ignore size))
  (if (env-exports env) (error "export-section only allowed once, found a second one"))
  (let* ((nels (read-leb stream)))
    (setf (env-exports env)
      (iter (for i below nels)
        (collect `(:name ,(read-wasm-utf8-string stream) ,(read-export-tag stream) ,(read-leb stream)))))))

(defun read-func-locals (stream)
  (let* ((nels (read-leb stream)))
    (iter (for i below nels) (collect `(,(read-leb stream) ,(read-valtype stream))))))
(defun read-code-section (size stream env)
  (declare (ignore size))
  (if (env-codes env) (error "code-section only allowed once, found a second one"))
  (let* ((nels (read-leb stream)))
    (setf (env-codes env)
      (iter (for i below nels)
        (collect `(:size ,(read-leb stream) :locals ,(read-func-locals stream) :body ,(read-expr stream)))))))
(defun read-data-section (size stream env)
  (declare (ignore size))
  (if (env-data env) (error "data-section only allowed once, found a second one"))
  (let* ((nels (read-leb stream)))
    (setf (env-data env)
      (iter (for i below nels)
        (collect `(:memidx ,(read-leb stream) :offset ,(read-expr stream) :bytes ,(read-wasm-byte-vec stream))))
    )))
(defun read-start-section (size stream env)
  (declare (ignore size))
  (if (env-start env) (error "start-section only allowed once, found a second one"))
  (let* ((funcidx (read-leb stream)))
    (setf (env-start env) funcidx)))

(defun read-wasm (stream env)
  (expect-magic 4 #x6d736100 stream)
  (expect-magic 4 #x00000001 stream)
  (iter
    (if verbose (format t ";; section head file-position:#x~(~8,'0x~)~%" (file-position stream)))
    (for tag-num = (read-byte stream nil nil))
    (while tag-num)
    (for tag = (section-tag-from-byte tag-num))
    (for size = (read-leb stream))
    (if verbose (format t ";; ~(tag:~s size:~d~)~%" tag size))
    (ecase tag ;; todo avoid case with some trickery with symbol-name, recast as symbol-function and funcall?
      (custom-section (read-custom-section size stream env))
      (type-section (read-type-section size stream env))
      (import-section (read-import-section size stream env))
      (function-section (read-function-section size stream env))
      (memory-section (read-memory-section size stream env))
      (global-section (read-global-section size stream env))
      (export-section (read-export-section size stream env))
      (code-section (read-code-section size stream env))
      (data-section (read-data-section size stream env))
      (start-section (read-start-section size stream env))
      )
    ;(format t "~(~s~)~%" env)
    )
  (if verbose (format t "~(~s~)~%" env))
  env)

(defun main (fname)
  (init-instr-lookup)
  (with-open-file (stream fname :element-type '(unsigned-byte 8))
    (read-wasm stream (make-env))))

(if verbose (format t "~(~s~)~%" (main "interp.wasm")))