Hot loading C wasm into the browser while preserving the state of the heap

core.cljs

;; on initial compile we want to make sure loader is loaded before reloader
;; these import orders ensure that. Just import this core ns from mainline.
(ns myproject.wasm.core
  (:require [myproject.wasm.heap]
            [myproject.wasm.loader]
            [myproject.wasm.reloader]))

example.c

typedef unsigned char uint8;
typedef unsigned int uint32;

__attribute__((export_name("calc")))
int calc(int a, int b)
{
  return a*b;
}

__attribute__((export_name("get_byte")))
uint8 get_byte(uint8 *buffer, int i)
{
  return buffer[i];
}

__attribute__((export_name("set_byte")))
void set_byte(uint8 *buffer, int i, uint8 v)
{
  buffer[i]=v;
}

heap.cljs

(ns myproject.wasm.heap)

(def page-size (* 64 1024))

(def initial-pages 16) ;; 1 MiB
(def maximum-pages (* 16 32)) ;; 32 MiB

;; a heap that gets reused between reloaded wasm code
(defonce memory
  (js/WebAssembly.Memory.
   (clj->js {:initial initial-pages
             :maximum maximum-pages})))

;; all the different views into the wasm heap
(defonce heap-uint8 (js/Uint8Array. (.-buffer memory)))
(defonce heap-uint16 (js/Uint16Array. (.-buffer memory)))
(defonce heap-uint32 (js/Uint32Array. (.-buffer memory)))
(defonce heap-int8 (js/Int8Array. (.-buffer memory)))
(defonce heap-int16 (js/Int16Array. (.-buffer memory)))
(defonce heap-int32 (js/Int32Array. (.-buffer memory)))

;; our basic heap allocator:
;; memory map is composed of [start-byte end-byte size] records.
;; memory is allocated out of first free sufficient space
;; gaps are bridged on freeing

(defrecord block [start-byte end-byte size])

(defonce memory-map
  (atom nil))

(defn init-memory-map! [heap-base]
  (swap! memory-map
         (fn [m]
           (or m (let [start heap-base
                       s (size)
                       end (dec s)
                       block (->block. start end s)
                       ]
                   {:free #{block}
                    :by-start {0 block}
                    :by-end {end block}
                    :alloc {}
                    :last-alloc nil})))))

(defn round-up-to-4-byte-boundary [b]
  (+ (* 4 (int (/ b 4)))
     (if (zero? (mod b 4)) 0 4)))

(defn allocate [{:keys [free by-start by-end alloc] :as allocated} bytes]
  ;; find first block of suitable size
  (let [aligned-bytes (round-up-to-4-byte-boundary bytes)
        block (->> free
                   (sort-by :start-byte)
                   (filter #(< bytes (:size %)))
                   first)]
    ;; if block is nil, we are out of memory and should grow the heap.
    (when block
      ;; slice the size of the front of the block
      (let [{:keys [start-byte end-byte size]} block
            new-start (+ start-byte aligned-bytes)
            new-block (->block. new-start end-byte (- size aligned-bytes))
            new-alloc (->block. start-byte (dec new-start) aligned-bytes)]
        {:free (-> free (disj block) (conj new-block))
         :by-start (-> by-start (dissoc start-byte) (assoc new-start new-block))
         :by-end (-> by-end (dissoc end-byte) (assoc end-byte new-block))
         :alloc (assoc alloc start-byte new-alloc)
         :last-alloc new-alloc}))))

(defn deallocate [{:keys [free by-start by-end  alloc last-alloc]} ptr]
  (let [{:keys [start-byte end-byte size] :as block} (alloc ptr)
        prev (by-end (dec start-byte))
        next (by-start (inc end-byte))
        ]
    (cond
      ;; bridges gap between two free blocks
      (and prev next)
      (let [new-block (->block. (:start-byte prev)
                                (:end-byte next)
                                (inc (- (:end-byte next) (:start-byte prev))))]
        {:free (-> free (disj prev next) (conj new-block))
         :by-start (-> by-start (assoc (:start-byte prev) new-block))
         :by-end (-> by-end (dissoc (:end-byte prev)) (assoc (:end-byte next) new-block))
         :alloc (dissoc alloc start-byte)
         :last-alloc last-alloc})

      ;; bridges gap to previous block
      prev
      (let [new-block (->block. (:start-byte prev)
                                end-byte
                                (inc (- end-byte (:start-byte prev))))]
        {:free (-> free (disj prev) (conj new-block))
         :by-start (-> by-start (assoc (:start-byte prev) new-block))
         :by-end (-> by-end (dissoc (:end-byte prev)) (assoc end-byte new-block))
         :alloc (dissoc alloc start-byte)
         :last-alloc last-alloc})

      ;; bridges with next block
      next
      (let [new-block (->block. start-byte
                                (:end-byte next)
                                (inc (- (:end-byte next) prev)))]
        {:free (-> free (disj next) (conj new-block))
         :by-start (-> by-start (dissoc (:start-byte next)) (assoc start-byte new-block))
         :by-end (-> by-end (assoc (:end-byte next) new-block))
         :alloc (dissoc alloc start-byte)
         :last-alloc last-alloc})

      ;; nothing joins. just remove it
      :else
      {:free (conj free block)
       :by-start (assoc by-start start-byte block)
       :by-end (assoc by-end end-byte block)
       :alloc (dissoc alloc start-byte)
       :last-alloc last-alloc})))

(defn free-space [{:keys [free]}]
  (reduce + (map :size free)))

(defn malloc [size]
  (let [start (:start-byte (:last-alloc (swap! memory-map allocate size)))]
    {:buffer (.subarray heap-uint8 start (+ start size))
     :address start}))

(defn free [{:keys [address]}]
  (swap! memory-map deallocate address))

loader.cljs

(ns myproject.wasm.loader
  (:require [myproject.wasm.heap :as heap]
            [cljs.core.async :refer [go]]
            [cljs.core.async.interop :refer [<p!]]))

;; load and reload compiled wasm modules keeping memory between them
(defonce module (atom {}))

(defn load-streaming [{:keys [name url]}]
  (go
    (let [result
          (<p!
           (js/WebAssembly.instantiateStreaming
            (js/fetch url)
            (clj->js {:env {:memory heap/memory}})))
          instance (.-instance result)
          exports (.-exports instance)
          heap-base (.-value (aget exports "__heap_base"))]
      (swap! module assoc name
             {:module (.-module result)
              :instance (.-instance result)
              :exports exports})
      (heap/init-memory-map! heap-base)
      (js/console.log (str "wasm " name " reloaded from " url " with heap-base " heap-base)))))

Makefile

CLANG_DIR = $(HOME)/clang+llvm-13.0.1-x86_64-linux-gnu-ubuntu-18.04
CLANG = $(CLANG_DIR)/bin/clang
LLC = $(CLANG_DIR)/bin/llc
LD = $(CLANG_DIR)/bin/wasm-ld

C_SRC := $(wildcard src/c/*.c)
OBJ_SRC := $(patsubst %.c, %.o, $(C_SRC))

%.o: %.c # delete competing implicit rule

%.ll: %.c
    $(CLANG) \
        --target=wasm32	\
        -emit-llvm \
        -c \
        -S \
        -std=c99 \
        -o $@ \
        -nostdlib \
        $<

%.o: %.ll
    $(LLC) \
        -march=wasm32 \
        -filetype=obj \
        $<

resources/public/wasm/mymodule.wasm: $(OBJ_SRC)
    $(LD) \
        --no-entry \
        --strip-all \
        --import-memory \
        --export=__heap_base \
        -o $@ \
        $^

reloader.cljs

(ns myproject.wasm.reloader)

(myproject.wasm.loader/load-streaming
 {:name :mymodule
  :url "wasm/mymodule.wasm"})

watch.clj

(require '[babashka.pods :as pods])
(pods/load-pod 'org.babashka/fswatcher "0.0.2")

(require '[pod.babashka.fswatcher :as fw]
         '[babashka.process :as p]
         '[clojure.string :as string])

(def path "src/c")
(def event-types #{:write})
(def path-extensions #{"c" "h"})
(def build-command "make resources/public/wasm/mymodule.wasm")
(def reload-hook-command "touch src/cljs/myproject/wasm/reloader.cljs")

(defn file-extension [s]
  (last (string/split s #"\.")))

(def watcher
  (fw/watch path
   (fn [{:keys [type path]}]
     (when (and (event-types type)
                (path-extensions (file-extension path)))
       (println "compiling...")
       (if (-> (p/process build-command {:inherit true})
               deref
               :exit
               zero?
               not)
         (println "failed!")
         (do
           (println "done.")
           @(p/process reload-hook-command)))))))

(.join (Thread/currentThread))

I came here from your blog-post @retrogradeorbit and thought you should have a link-back: https://epiccastle.io/blog/hot-loading-wasm/
Well done btw!