Custom Brotli Dictionaries

Dockerfile

@@ -46,11 +46,15 @@ RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --de
 COPY ./Makefile ./
 COPY arbitrator/Cargo.* arbitrator/
 COPY arbitrator/arbutil arbitrator/arbutil
+COPY arbitrator/brotli arbitrator/brotli
 COPY arbitrator/caller-env arbitrator/caller-env
 COPY arbitrator/prover arbitrator/prover
 COPY arbitrator/wasm-libraries arbitrator/wasm-libraries
 COPY arbitrator/tools/wasmer arbitrator/tools/wasmer
+COPY brotli brotli
+COPY scripts/build-brotli.sh scripts/
 COPY --from=brotli-wasm-export / target/
+RUN apt-get update && apt-get install -y cmake
 RUN . ~/.cargo/env && NITRO_BUILD_IGNORE_TIMESTAMPS=1 RUSTFLAGS='-C symbol-mangling-version=v0' make build-wasm-libs
 
 FROM scratch as wasm-libs-export
@@ -93,12 +97,17 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
 COPY arbitrator/Cargo.* arbitrator/
 COPY ./Makefile ./
 COPY arbitrator/arbutil arbitrator/arbutil
+COPY arbitrator/brotli arbitrator/brotli
 COPY arbitrator/caller-env arbitrator/caller-env
 COPY arbitrator/prover arbitrator/prover
 COPY arbitrator/wasm-libraries arbitrator/wasm-libraries
 COPY arbitrator/jit arbitrator/jit
 COPY arbitrator/stylus arbitrator/stylus
 COPY arbitrator/tools/wasmer arbitrator/tools/wasmer
+COPY --from=brotli-wasm-export / target/
+COPY scripts/build-brotli.sh scripts/
+COPY brotli brotli
+RUN apt-get update && apt-get install -y cmake
 RUN NITRO_BUILD_IGNORE_TIMESTAMPS=1 make build-prover-header
 
 FROM scratch as prover-header-export
@@ -113,8 +122,10 @@ RUN wget -O - https://apt.llvm.org/llvm-snapshot.gpg.key | apt-key add - && \
     add-apt-repository 'deb http://apt.llvm.org/bullseye/ llvm-toolchain-bullseye-15 main' && \
     apt-get update && \
     apt-get install -y llvm-15-dev libclang-common-15-dev libpolly-15-dev
+COPY --from=brotli-library-export / target/
 COPY arbitrator/Cargo.* arbitrator/
 COPY arbitrator/arbutil arbitrator/arbutil
+COPY arbitrator/brotli arbitrator/brotli
 COPY arbitrator/caller-env arbitrator/caller-env
 COPY arbitrator/prover/Cargo.toml arbitrator/prover/
 COPY arbitrator/jit/Cargo.toml arbitrator/jit/
@@ -134,7 +145,9 @@ COPY arbitrator/prover arbitrator/prover
 COPY arbitrator/wasm-libraries arbitrator/wasm-libraries
 COPY arbitrator/jit arbitrator/jit
 COPY arbitrator/stylus arbitrator/stylus
-COPY --from=brotli-library-export / target/
+COPY --from=brotli-wasm-export / target/
+COPY scripts/build-brotli.sh scripts/
+COPY brotli brotli
 RUN touch -a -m arbitrator/prover/src/lib.rs
 RUN NITRO_BUILD_IGNORE_TIMESTAMPS=1 make build-prover-lib
 RUN NITRO_BUILD_IGNORE_TIMESTAMPS=1 make build-prover-bin
@@ -156,6 +169,7 @@ COPY --from=wasm-libs-builder /workspace/arbitrator/prover/ arbitrator/prover/
 COPY --from=wasm-libs-builder /workspace/arbitrator/tools/wasmer/ arbitrator/tools/wasmer/
 COPY --from=wasm-libs-builder /workspace/arbitrator/wasm-libraries/ arbitrator/wasm-libraries/
 COPY --from=wasm-libs-builder /workspace/arbitrator/arbutil arbitrator/arbutil
+COPY --from=wasm-libs-builder /workspace/arbitrator/brotli arbitrator/brotli
 COPY --from=wasm-libs-builder /workspace/arbitrator/caller-env arbitrator/caller-env
 COPY --from=wasm-libs-builder /workspace/.make/ .make/
 COPY ./Makefile ./

Makefile

@@ -28,7 +28,7 @@ ifneq ($(origin NITRO_MODIFIED),undefined)
 endif
 
 ifneq ($(origin GOLANG_LDFLAGS),undefined)
- GOLANG_PARAMS = -ldflags="$(GOLANG_LDFLAGS)"
+ GOLANG_PARAMS = -ldflags="-extldflags '-ldl' $(GOLANG_LDFLAGS)"
 endif
 
 precompile_names = AddressTable Aggregator BLS Debug FunctionTable GasInfo Info osTest Owner RetryableTx Statistics Sys
@@ -37,7 +37,7 @@ precompiles = $(patsubst %,./solgen/generated/%.go, $(precompile_names))
 output_root=target
 output_latest=$(output_root)/machines/latest
 
-repo_dirs = arbos arbnode arbutil arbstate cmd das precompiles solgen system_tests util validator wavmio
+repo_dirs = arbos arbcompress arbnode arbutil arbstate cmd das precompiles solgen system_tests util validator wavmio
 go_source.go = $(wildcard $(patsubst %,%/*.go, $(repo_dirs)) $(patsubst %,%/*/*.go, $(repo_dirs)))
 go_source.s  = $(wildcard $(patsubst %,%/*.s, $(repo_dirs)) $(patsubst %,%/*/*.s, $(repo_dirs)))
 go_source = $(go_source.go) $(go_source.s)
@@ -47,12 +47,12 @@ color_reset = "\e[0;0m"
 
 done = "%bdone!%b\n" $(color_pink) $(color_reset)
 
-replay_deps=arbos wavmio arbstate arbcompress solgen/go/node-interfacegen blsSignatures cmd/replay
-
 replay_wasm=$(output_latest)/replay.wasm
 
+arb_brotli_files = $(wildcard arbitrator/brotli/src/*.* arbitrator/brotli/src/*/*.* arbitrator/brotli/*.toml arbitrator/brotli/*.rs) .make/cbrotli-lib .make/cbrotli-wasm
+
 arbitrator_generated_header=$(output_root)/include/arbitrator.h
-arbitrator_wasm_libs=$(patsubst %, $(output_root)/machines/latest/%.wasm, forward wasi_stub host_io soft-float brotli user_host program_exec)
+arbitrator_wasm_libs=$(patsubst %, $(output_root)/machines/latest/%.wasm, forward wasi_stub host_io soft-float arbcompress user_host program_exec)
 arbitrator_stylus_lib=$(output_root)/lib/libstylus.a
 prover_bin=$(output_root)/bin/prover
 arbitrator_jit=$(output_root)/bin/jit
@@ -74,14 +74,14 @@ WASI_SYSROOT?=/opt/wasi-sdk/wasi-sysroot
 
 arbitrator_wasm_lib_flags=$(patsubst %, -l %, $(arbitrator_wasm_libs))
 
-rust_arbutil_files = $(wildcard arbitrator/arbutil/src/*.* arbitrator/arbutil/src/*/*.* arbitrator/arbutil/*.toml arbitrator/caller-env/src/*.* arbitrator/caller-env/src/*/*.* arbitrator/caller-env/*.toml)
+rust_arbutil_files = $(wildcard arbitrator/arbutil/src/*.* arbitrator/arbutil/src/*/*.* arbitrator/arbutil/*.toml arbitrator/caller-env/src/*.* arbitrator/caller-env/src/*/*.* arbitrator/caller-env/*.toml) .make/cbrotli-lib
 
 prover_direct_includes = $(patsubst %,$(output_latest)/%.wasm, forward forward_stub)
-prover_src = arbitrator/prover/src
-rust_prover_files = $(wildcard $(prover_src)/*.* $(prover_src)/*/*.* arbitrator/prover/*.toml) $(rust_arbutil_files) $(prover_direct_includes)
+prover_dir = arbitrator/prover/
+rust_prover_files = $(wildcard $(prover_dir)/src/*.* $(prover_dir)/src/*/*.* $(prover_dir)/*.toml $(prover_dir)/*.rs) $(rust_arbutil_files) $(prover_direct_includes) $(arb_brotli_files)
 
 wasm_lib = arbitrator/wasm-libraries
-wasm_lib_deps = $(wildcard $(wasm_lib)/$(1)/*.toml $(wasm_lib)/$(1)/src/*.rs $(wasm_lib)/$(1)/*.rs) $(rust_arbutil_files) .make/machines
+wasm_lib_deps = $(wildcard $(wasm_lib)/$(1)/*.toml $(wasm_lib)/$(1)/src/*.rs $(wasm_lib)/$(1)/*.rs) $(rust_arbutil_files) $(arb_brotli_files) .make/machines
 wasm_lib_go_abi = $(call wasm_lib_deps,go-abi)
 wasm_lib_forward = $(call wasm_lib_deps,forward)
 wasm_lib_user_host_trait = $(call wasm_lib_deps,user-host-trait)
@@ -274,7 +274,7 @@ $(arbitrator_stylus_lib): $(DEP_PREDICATE) $(stylus_files)
 	cargo build --manifest-path arbitrator/Cargo.toml --release --lib -p stylus ${CARGOFLAGS}
 	install arbitrator/target/release/libstylus.a $@
 
-$(arbitrator_jit): $(DEP_PREDICATE) .make/cbrotli-lib $(jit_files)
+$(arbitrator_jit): $(DEP_PREDICATE) $(jit_files)
 	mkdir -p `dirname $(arbitrator_jit)`
 	cargo build --manifest-path arbitrator/Cargo.toml --release -p jit ${CARGOFLAGS}
 	install arbitrator/target/release/jit $@
@@ -349,9 +349,9 @@ $(output_latest)/user_test.wasm: $(DEP_PREDICATE) $(call wasm_lib_deps,user-test
 	cargo build --manifest-path arbitrator/wasm-libraries/Cargo.toml --release --target wasm32-wasi --package user-test
 	install arbitrator/wasm-libraries/$(wasm32_wasi)/user_test.wasm $@
 
-$(output_latest)/brotli.wasm: $(DEP_PREDICATE) $(call wasm_lib_deps,brotli) $(wasm_lib_go_abi) .make/cbrotli-wasm
-	cargo build --manifest-path arbitrator/wasm-libraries/Cargo.toml --release --target wasm32-wasi --package brotli
-	install arbitrator/wasm-libraries/$(wasm32_wasi)/brotli.wasm $@
+$(output_latest)/arbcompress.wasm: $(DEP_PREDICATE) $(call wasm_lib_deps,brotli) $(wasm_lib_go_abi)
+	cargo build --manifest-path arbitrator/wasm-libraries/Cargo.toml --release --target wasm32-wasi --package arbcompress
+	install arbitrator/wasm-libraries/$(wasm32_wasi)/arbcompress.wasm $@
 
 $(output_latest)/forward.wasm: $(DEP_PREDICATE) $(wasm_lib_forward) .make/machines
 	cargo run --manifest-path $(forward_dir)/Cargo.toml -- --path $(forward_dir)/forward.wat
@@ -363,7 +363,7 @@ $(output_latest)/forward_stub.wasm: $(DEP_PREDICATE) $(wasm_lib_forward) .make/m
 
 $(output_latest)/machine.wavm.br: $(DEP_PREDICATE) $(prover_bin) $(arbitrator_wasm_libs) $(replay_wasm)
 	$(prover_bin) $(replay_wasm) --generate-binaries $(output_latest) \
-	$(patsubst %,-l $(output_latest)/%.wasm, forward soft-float wasi_stub host_io user_host brotli program_exec)
+	$(patsubst %,-l $(output_latest)/%.wasm, forward soft-float wasi_stub host_io user_host arbcompress program_exec)
 
 $(arbitrator_cases)/%.wasm: $(arbitrator_cases)/%.wat
 	wat2wasm $< -o $@

arbcompress/compress_common.go

@@ -1,13 +1,13 @@
-// Copyright 2021-2022, Offchain Labs, Inc.
+// Copyright 2021-2024, Offchain Labs, Inc.
 // For license information, see https://github.com/nitro/blob/master/LICENSE
 
 package arbcompress
 
-type BrotliStatus = uint32
+type Dictionary uint32
 
 const (
-	BrotliFailure uint32 = iota
-	BrotliSuccess
+	EmptyDictionary Dictionary = iota
+	StylusProgramDictionary
 )
 
 const LEVEL_FAST = 0
@@ -19,5 +19,5 @@ func compressedBufferSizeFor(length int) int {
 }
 
 func CompressFast(input []byte) ([]byte, error) {
-	return compressLevel(input, LEVEL_FAST)
+	return Compress(input, LEVEL_FAST, EmptyDictionary)
 }

arbcompress/native.go

@@ -0,0 +1,76 @@
+// Copyright 2021-2022, Offchain Labs, Inc.
+// For license information, see https://github.com/nitro/blob/master/LICENSE
+
+//go:build !wasm
+// +build !wasm
+
+package arbcompress
+
+/*
+#cgo CFLAGS: -g -Wall -I${SRCDIR}/../target/include/
+#cgo LDFLAGS: ${SRCDIR}/../target/lib/libstylus.a -lm
+#include "arbitrator.h"
+*/
+import "C"
+import "fmt"
+
+type u8 = C.uint8_t
+type u32 = C.uint32_t
+type usize = C.size_t
+
+type brotliBool = uint32
+type brotliBuffer = C.BrotliBuffer
+
+const (
+	brotliFalse brotliBool = iota
+	brotliTrue
+)
+
+func CompressWell(input []byte) ([]byte, error) {
+	return Compress(input, LEVEL_WELL, EmptyDictionary)
+}
+
+func Compress(input []byte, level int, dictionary Dictionary) ([]byte, error) {
+	maxSize := compressedBufferSizeFor(len(input))
+	output := make([]byte, maxSize)
+	outbuf := sliceToBuffer(output)
+	inbuf := sliceToBuffer(input)
+
+	status := C.brotli_compress(inbuf, outbuf, C.Dictionary(dictionary), u32(level))
+	if status != C.BrotliStatus_Success {
+		return nil, fmt.Errorf("failed decompression: %d", status)
+	}
+	output = output[:*outbuf.len]
+	return output, nil
+}
+
+func Decompress(input []byte, maxSize int) ([]byte, error) {
+	return DecompressWithDictionary(input, maxSize, EmptyDictionary)
+}
+
+func DecompressWithDictionary(input []byte, maxSize int, dictionary Dictionary) ([]byte, error) {
+	output := make([]byte, maxSize)
+	outbuf := sliceToBuffer(output)
+	inbuf := sliceToBuffer(input)
+
+	status := C.brotli_decompress(inbuf, outbuf, C.Dictionary(dictionary))
+	if status != C.BrotliStatus_Success {
+		return nil, fmt.Errorf("failed decompression: %d", status)
+	}
+	if *outbuf.len > usize(maxSize) {
+		return nil, fmt.Errorf("failed decompression: result too large: %d", *outbuf.len)
+	}
+	output = output[:*outbuf.len]
+	return output, nil
+}
+
+func sliceToBuffer(slice []byte) brotliBuffer {
+	count := usize(len(slice))
+	if count == 0 {
+		slice = []byte{0x00} // ensures pointer is not null (shouldn't be necessary, but brotli docs are picky about NULL)
+	}
+	return brotliBuffer{
+		ptr: (*u8)(&slice[0]),
+		len: &count,
+	}
+}

arbcompress/compress_wasm.go → arbcompress/wasm.go

@@ -13,25 +13,20 @@ import (
 	"github.com/offchainlabs/nitro/arbutil"
 )
 
+type brotliStatus = uint32
+
+const (
+	brotliFailure brotliStatus = iota
+	brotliSuccess
+)
+
 //go:wasmimport arbcompress brotli_compress
-func brotliCompress(inBuf unsafe.Pointer, inBufLen uint32, outBuf unsafe.Pointer, outBufLen unsafe.Pointer, level, windowSize uint32) BrotliStatus
+func brotliCompress(inBuf unsafe.Pointer, inLen uint32, outBuf unsafe.Pointer, outLen unsafe.Pointer, level, windowSize uint32, dictionary Dictionary) brotliStatus
 
 //go:wasmimport arbcompress brotli_decompress
-func brotliDecompress(inBuf unsafe.Pointer, inBufLen uint32, outBuf unsafe.Pointer, outBufLen unsafe.Pointer) BrotliStatus
-
-func Decompress(input []byte, maxSize int) ([]byte, error) {
-	outBuf := make([]byte, maxSize)
-	outLen := uint32(len(outBuf))
-	status := brotliDecompress(
-		arbutil.SliceToUnsafePointer(input), uint32(len(input)), arbutil.SliceToUnsafePointer(outBuf), unsafe.Pointer(&outLen),
-	)
-	if status != BrotliSuccess {
-		return nil, fmt.Errorf("failed decompression")
-	}
-	return outBuf[:outLen], nil
-}
+func brotliDecompress(inBuf unsafe.Pointer, inLen uint32, outBuf unsafe.Pointer, outLen unsafe.Pointer, dictionary Dictionary) brotliStatus
 
-func compressLevel(input []byte, level uint32) ([]byte, error) {
+func Compress(input []byte, level uint32, dictionary Dictionary) ([]byte, error) {
 	maxOutSize := compressedBufferSizeFor(len(input))
 	outBuf := make([]byte, maxOutSize)
 	outLen := uint32(len(outBuf))
@@ -40,9 +35,30 @@ func compressLevel(input []byte, level uint32) ([]byte, error) {
 		arbutil.SliceToUnsafePointer(outBuf), unsafe.Pointer(&outLen),
 		level,
 		WINDOW_SIZE,
+		dictionary,
 	)
-	if status != BrotliSuccess {
+	if status != brotliSuccess {
 		return nil, fmt.Errorf("failed compression")
 	}
 	return outBuf[:outLen], nil
 }
+
+func Decompress(input []byte, maxSize int) ([]byte, error) {
+	return DecompressWithDictionary(input, maxSize, EmptyDictionary)
+}
+
+func DecompressWithDictionary(input []byte, maxSize int, dictionary Dictionary) ([]byte, error) {
+	outBuf := make([]byte, maxSize)
+	outLen := uint32(len(outBuf))
+	status := brotliDecompress(
+		arbutil.SliceToUnsafePointer(input),
+		uint32(len(input)),
+		arbutil.SliceToUnsafePointer(outBuf),
+		unsafe.Pointer(&outLen),
+		dictionary,
+	)
+	if status != brotliSuccess {
+		return nil, fmt.Errorf("failed decompression")
+	}
+	return outBuf[:outLen], nil
+}