Skip to content

Instantly share code, notes, and snippets.

@CryZe

CryZe/aoc_2024_17.rs

Last active December 30, 2024 19:03
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save CryZe/41341ff49fc9f41c4fed5345dc6ebb67 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save CryZe/41341ff49fc9f41c4fed5345dc6ebb67 to your computer and use it in GitHub Desktop.
Download ZIP
Advent of Code 2024 Day 17 WebAssembly Compiler
Raw
aoc_2024_17.rs
use std::fs;
use gimli::{
write::{
Address, AttributeValue, DwarfUnit, EndianVec, Expression, LineProgram, LineString, Range,
RangeList, Sections, UnitEntryId,
},
Encoding, Format, LineEncoding, LittleEndian,
};
use wasm_encoder::{
BlockType, CodeSection, ConstExpr, CustomSection, EntityType, ExportKind, ExportSection,
Function, FunctionSection, GlobalSection, GlobalType, ImportSection, IndirectNameMap,
Instruction, Module, NameMap, NameSection, TypeSection, ValType,
};
fn parse(rem: &str) -> Ast {
let rem = rem.strip_prefix("Register A: ").unwrap();
let (a, rem) = rem.split_once('\n').unwrap();
let rem = rem.strip_prefix("Register B: ").unwrap();
let (b, rem) = rem.split_once('\n').unwrap();
let rem = rem.strip_prefix("Register C: ").unwrap();
let (c, rem) = rem.split_once('\n').unwrap();
let mut rem = rem.strip_prefix("\nProgram: ").unwrap().as_bytes();
let mut instructions = Vec::new();
loop {
match rem {
[ins, b',', op, next, new_rem @ ..] => {
instructions.push(match ins {
b'0' => Ins::Adv(combo(*op)),
b'1' => Ins::Bxl(*op - b'0'),
b'2' => Ins::Bst(combo(*op)),
b'3' => Ins::Jnz(*op - b'0'),
b'4' => Ins::Bxc,
b'5' => Ins::Out(combo(*op)),
b'6' => Ins::Bdv(combo(*op)),
b'7' => Ins::Cdv(combo(*op)),
_ => unreachable!(),
});
rem = new_rem;
if *next == b'\n' {
break;
}
}
_ => unreachable!(),
}
}
Ast {
a: a.parse().unwrap(),
b: b.parse().unwrap(),
c: c.parse().unwrap(),
instructions,
}
}
fn combo(b: u8) -> Combo {
match b {
b'0' => Combo::Zero,
b'1' => Combo::One,
b'2' => Combo::Two,
b'3' => Combo::Three,
b'4' => Combo::A,
b'5' => Combo::B,
b'6' => Combo::C,
_ => unreachable!(),
}
}
#[derive(Debug)]
struct Ast {
a: i32,
b: i32,
c: i32,
instructions: Vec<Ins>,
}
#[derive(Debug)]
enum Combo {
Zero,
One,
Two,
Three,
A,
B,
C,
}
#[derive(Debug)]
enum Ins {
Adv(Combo),
Bxl(u8),
Bst(Combo),
Jnz(u8),
Bxc,
Out(Combo),
Bdv(Combo),
Cdv(Combo),
}
fn encode(ast: &Ast) -> Vec<u8> {
let a = 0;
let b = 1;
let c = 2;
let ip = 3;
let mut module = Module::new();
let mut types = TypeSection::new();
types.ty().function([ValType::I32; 3], []);
types.ty().function([ValType::I32; 1], []);
let mut imports = ImportSection::new();
imports.import("env", "out", EntityType::Function(1));
let mut functions = FunctionSection::new();
functions.function(0);
let mut globals = GlobalSection::new();
globals.global(
GlobalType {
val_type: ValType::I32,
mutable: false,
shared: false,
},
&ConstExpr::i32_const(ast.a),
);
globals.global(
GlobalType {
val_type: ValType::I32,
mutable: false,
shared: false,
},
&ConstExpr::i32_const(ast.b),
);
globals.global(
GlobalType {
val_type: ValType::I32,
mutable: false,
shared: false,
},
&ConstExpr::i32_const(ast.c),
);
let mut exports = ExportSection::new();
exports.export("program", ExportKind::Func, 1);
exports.export("A", ExportKind::Global, 0);
exports.export("B", ExportKind::Global, 1);
exports.export("C", ExportKind::Global, 2);
let mut codes = CodeSection::new();
let mut f = Function::new([(1, ValType::I32)]);
let mut jump_locations = ast
.instructions
.iter()
.filter_map(|ins| match ins {
Ins::Jnz(operand) if (*operand as usize) < ast.instructions.len() => {
Some(*operand as u32)
}
_ => None,
})
.collect::<Vec<_>>();
jump_locations.sort_unstable();
jump_locations.dedup();
f.instruction(&Instruction::Loop(BlockType::Empty));
for _ in 0..jump_locations.len() as u32 {
f.instruction(&Instruction::Block(BlockType::Empty));
}
f.instruction(&Instruction::LocalGet(3))
.instruction(&Instruction::BrTable(
(0..jump_locations.len() as u32).collect::<Vec<_>>().into(),
0,
));
let mut rem_jump_locations = &jump_locations[..];
let mut offsets = Vec::new();
for (cur_pos, instruction) in ast.instructions.iter().enumerate() {
if rem_jump_locations
.first()
.is_some_and(|&loc| loc == 2 * cur_pos as u32)
{
rem_jump_locations = &rem_jump_locations[1..];
f.instruction(&Instruction::End);
}
offsets.push(f.byte_len() + 1);
match instruction {
Ins::Adv(combo) => {
f.instruction(&Instruction::LocalGet(a))
.instruction(&lower_combo(combo))
.instruction(&Instruction::I32ShrU)
.instruction(&Instruction::LocalSet(a));
}
Ins::Bdv(combo) => {
f.instruction(&Instruction::LocalGet(a))
.instruction(&lower_combo(combo))
.instruction(&Instruction::I32ShrU)
.instruction(&Instruction::LocalSet(b));
}
Ins::Cdv(combo) => {
f.instruction(&Instruction::LocalGet(a))
.instruction(&lower_combo(combo))
.instruction(&Instruction::I32ShrU)
.instruction(&Instruction::LocalSet(c));
}
Ins::Bxl(operand) => {
f.instruction(&Instruction::LocalGet(b))
.instruction(&Instruction::I32Const(*operand as i32))
.instruction(&Instruction::I32Xor)
.instruction(&Instruction::LocalSet(b));
}
Ins::Bst(combo) => {
f.instruction(&lower_combo(combo))
.instruction(&Instruction::I32Const(0b111))
.instruction(&Instruction::I32And)
.instruction(&Instruction::LocalSet(b));
}
Ins::Jnz(pos) => {
f.instruction(&Instruction::LocalGet(a))
.instruction(&Instruction::If(BlockType::Empty));
if let Some(label_idx) = jump_locations.iter().position(|&x| x == *pos as u32) {
f.instruction(&Instruction::I32Const(label_idx as i32))
.instruction(&Instruction::LocalSet(ip))
.instruction(&Instruction::Br(rem_jump_locations.len() as u32 + 1));
} else {
f.instruction(&Instruction::Return);
}
f.instruction(&Instruction::End);
}
Ins::Bxc => {
f.instruction(&Instruction::LocalGet(b))
.instruction(&Instruction::LocalGet(c))
.instruction(&Instruction::I32Xor)
.instruction(&Instruction::LocalSet(b));
}
Ins::Out(combo) => {
f.instruction(&lower_combo(combo))
.instruction(&Instruction::I32Const(0b111))
.instruction(&Instruction::I32And)
.instruction(&Instruction::Call(0));
}
}
}
f.instruction(&Instruction::End);
f.instruction(&Instruction::End);
let main_size = f.byte_len();
codes.function(&f);
let mut names = NameSection::new();
let mut function_names = NameMap::new();
function_names.append(0, "out");
function_names.append(1, "program");
let mut locals = IndirectNameMap::new();
let mut foo_locals = NameMap::new();
foo_locals.append(a, "A");
foo_locals.append(b, "B");
foo_locals.append(c, "C");
foo_locals.append(ip, "IP");
locals.append(1, &foo_locals);
names.functions(&function_names);
names.locals(&locals);
module.section(&types);
module.section(&imports);
module.section(&functions);
module.section(&globals);
module.section(&exports);
module.section(&codes);
module.section(&names);
dwarf(&mut module, &offsets, main_size as u64);
module.finish()
}
fn dwarf(module: &mut Module, offsets: &[usize], main_size: u64) {
let comp_dir = *br"P:\advent-of-code-2024-17\src";
let file_name = *b"input.txt";
let main_name = *b"program";
let main_address = Address::Constant(0);
let encoding = Encoding {
format: Format::Dwarf32,
version: 4,
address_size: 4,
};
let mut dwarf = DwarfUnit::new(encoding);
let range_list_id = dwarf.unit.ranges.add(RangeList(vec![Range::StartLength {
begin: main_address,
length: main_size,
}]));
let root = dwarf.unit.root();
let entry = dwarf.unit.get_mut(root);
entry.set(
gimli::DW_AT_producer,
AttributeValue::String((*b"day17").into()),
);
entry.set(
gimli::DW_AT_language,
AttributeValue::Language(gimli::DW_LANG_C11),
);
entry.set(gimli::DW_AT_name, AttributeValue::String(file_name.into()));
entry.set(
gimli::DW_AT_comp_dir,
AttributeValue::String(comp_dir.into()),
);
entry.set(gimli::DW_AT_low_pc, AttributeValue::Address(main_address));
entry.set(
gimli::DW_AT_ranges,
AttributeValue::RangeListRef(range_list_id),
);
let line_strings = &mut dwarf.line_strings;
let mut line_program = LineProgram::new(
encoding,
LineEncoding::default(),
LineString::new(comp_dir, encoding, line_strings),
LineString::new(file_name, encoding, line_strings),
None,
);
let dir_id = line_program.default_directory();
let file_string = LineString::new(file_name, encoding, line_strings);
let file_id = line_program.add_file(file_string, dir_id, None);
line_program.begin_sequence(Some(main_address));
let mut column = 10;
line_program.row().file = file_id;
line_program.row().line = 5;
line_program.row().column = 1;
line_program.generate_row();
for &offset in offsets {
line_program.row().address_offset = offset as u64;
line_program.row().column = column;
column += 4;
line_program.generate_row();
}
line_program.end_sequence(main_size);
dwarf.unit.line_program = line_program;
let subprogram = dwarf.unit.add(root, gimli::DW_TAG_subprogram);
let entry = dwarf.unit.get_mut(subprogram);
entry.set(gimli::DW_AT_external, AttributeValue::Flag(true));
entry.set(gimli::DW_AT_name, AttributeValue::String(main_name.into()));
entry.set(
gimli::DW_AT_decl_file,
AttributeValue::FileIndex(Some(file_id)),
);
entry.set(gimli::DW_AT_decl_line, AttributeValue::Udata(5));
entry.set(gimli::DW_AT_decl_column, AttributeValue::Udata(10));
entry.set(gimli::DW_AT_low_pc, AttributeValue::Address(main_address));
entry.set(gimli::DW_AT_high_pc, AttributeValue::Udata(main_size));
let u32_ty = mk_primitive(&mut dwarf, "u32", 4, gimli::DW_ATE_unsigned);
for (i, var_name) in ["A", "B", "C"].into_iter().enumerate() {
let param = dwarf.unit.add(subprogram, gimli::DW_TAG_variable);
let entry = dwarf.unit.get_mut(param);
entry.set(gimli::DW_AT_name, AttributeValue::String(var_name.into()));
let mut loc = Expression::new();
loc.op_wasm_local(i as _);
loc.op(gimli::DW_OP_stack_value);
entry.set(gimli::DW_AT_location, AttributeValue::Exprloc(loc));
entry.set(gimli::DW_AT_type, AttributeValue::UnitRef(u32_ty));
}
let mut sections = Sections::new(EndianVec::new(LittleEndian));
dwarf.write(&mut sections).unwrap();
sections
.for_each(|id, data| {
if data.slice().is_empty() {
return Ok::<(), ()>(());
}
module.section(&CustomSection {
name: id.name().into(),
data: data.slice().into(),
});
Ok::<(), ()>(())
})
.unwrap();
}
fn mk_primitive(
dwarf: &mut DwarfUnit,
name: &str,
byte_size: u8,
encoding: gimli::DwAte,
) -> UnitEntryId {
let name = dwarf.strings.add(name);
let root = dwarf.unit.root();
let ty = dwarf.unit.add(root, gimli::DW_TAG_base_type);
let entry = dwarf.unit.get_mut(ty);
entry.set(gimli::DW_AT_name, AttributeValue::StringRef(name));
entry.set(gimli::DW_AT_byte_size, AttributeValue::Data1(byte_size));
entry.set(gimli::DW_AT_encoding, AttributeValue::Encoding(encoding));
ty
}
fn lower_combo(combo: &Combo) -> Instruction<'_> {
match combo {
Combo::Zero => Instruction::I32Const(0),
Combo::One => Instruction::I32Const(1),
Combo::Two => Instruction::I32Const(2),
Combo::Three => Instruction::I32Const(3),
Combo::A => Instruction::LocalGet(0),
Combo::B => Instruction::LocalGet(1),
Combo::C => Instruction::LocalGet(2),
}
}
fn main() {
println!("Parsing...");
let ast = dbg!(parse(include_str!("input.txt")));
println!("Codegen...");
let wasm = encode(&ast);
println!("Writing wasm...");
fs::write("foo.wasm", &wasm).unwrap();
}
Raw
output.wat
(module
(type $0 (func (param i32)))
(type $1 (func (param i32 i32 i32)))
(import "env" "out" (func $out (param i32)))
(global $global$0 i32 (i32.const 30344604))
(global $global$1 i32 (i32.const 0))
(global $global$2 i32 (i32.const 0))
(export "program" (func $program))
(export "A" (global $global$0))
(export "B" (global $global$1))
(export "C" (global $global$2))
(func $program (param $A i32) (param $B i32) (param $C i32)
(local $IP i32)
(loop $label$1
(block $label$2
(br_table $label$2 $label$2
(local.get $IP)
)
)
(local.set $B
(i32.and
(local.get $A)
(i32.const 7)
)
)
(local.set $B
(i32.xor
(local.get $B)
(i32.const 1)
)
)
(local.set $C
(i32.shr_u
(local.get $A)
(local.get $B)
)
)
(local.set $B
(i32.xor
(local.get $B)
(i32.const 5)
)
)
(local.set $B
(i32.xor
(local.get $B)
(local.get $C)
)
)
(local.set $A
(i32.shr_u
(local.get $A)
(i32.const 3)
)
)
(call $out
(i32.and
(local.get $B)
(i32.const 7)
)
)
(if
(local.get $A)
(then
(local.set $IP
(i32.const 0)
)
(br $label$1)
)
)
)
)
;; custom section ".debug_abbrev", size 61
;; custom section ".debug_str", size 4
;; custom section ".debug_line", size 82
;; custom section ".debug_ranges", size 16
;; custom section ".debug_info", size 134
)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment