Nested declarative macros for enum dispatch in Rust

serde_arrow relies heavily on Rust macros to remove code duplication. One prominent example is the bytecode introduced in version 0.7. It compiles the underlying state machines into a list of instructions that are then interpreted. Each instruction is one of a large number of variants in an enum. In this post I would like to summarize how serde_arrow uses nested declarative macross, that is macros defining macros, to avoid code duplication in the interpreter.

But first, what is serde_arrow? It's a Rust package that allows to serialize and deserialize Rust objects from and to Arrow arrays. Think of converting a list of Rust structs to a data frame and back again. To do so it uses the machinery provided by serde.

Consider the following struct

#[derive(Serialize)]
struct Example {
    int_value: i32,
    float_value: f32,
}

The derive macro of serde generates roughly the following code for serialization

impl Serialize for Example {
    fn serialize<S: Serializer>(&self, ser: S) -> Result<S::Ok, S::Error> {
        let mut state = ser.serialize_struct("Example", 2)?;
        state.serialize_field("int_value", &self.int_value)?;
        state.serialize_field("float_value", &self.float_value)?;
        state.end()
    }
}

serde_arrow converts the serde calls into a list of events

[
    Event::StartStruct,
    Event::Str("int_value"),
    Event::I32(..),
    Event::Str("float_value"),
    Event::F32(..),
    Event::EndStruct,
]

that are then interpreted to build the arrays, in previous versions by a nested state machine. However, the overhead to drive the state machine can be considerable in particular for nested structures. To improve performance, newer versions of serde_arrow compile the state machine into a bytecode.

The bytecode program that corresponds to the Example struct looks similar to

(
    OuterSequenceStart {..},
    OuterSequenceItem {..},
    OuterRecordStart {..},
    OuterRecordField { 
        field_name: String::from("int_value"),
        ..,
    },
    PushI32 {..},
    OuterRecordField { 
        field_name: String::from("float_value"),
        ..,
    },
    PushF32 {..},
    OuterRecordEnd {..},
    OuterSequenceEnd {..},
    ProgramEnd {..},
)

In the full code, each instruction is wrapped into a common enum that is defined as

enum Bytecode {
    OuterSequenceStart(OuterSequenceStart),
    OuterSequenceEnd(OuterSequenceEnd),
    OuterSequenceItem(OuterSequenceItem),
    OuterRecordStart(OuterRecordStart),
    OuterRecordField(OuterRecordField), 
    OuterRecordEnd(OuterRecordEnd),
    ProgramEnd(ProgramEnd),
    PushI32(PushI32),
    PushF32(PushF32),
}

The interpreter for this program accepts individual events from serde and dispatches them to the current instruction:

impl Interpreter {
    fn accept(&mut self, event: &Event) -> Result<()> {
        match &self.program[self.current_instruction] {
            Bytecode::OuterSequenceStart(instr) => instr.accept(event),
            Bytecode::OuterSequenceEnd(instr) => instr.accept(event),
            Bytecode::OuterSequenceItem(instr) => instr.accept(event),
            /* ... */       
        }
    }
}

Note how the same code is repeated for each variant. To remove code duplication, serde_arrow uses a macro to define

  • the Bytecode enum,
  • the variant structs, and
  • a dispatch macro to apply a code block for each variant
macro_rules! define_bytecode {
    (
        $(
            $variant:ident {
                $( $field:ident: $ty:ty, )*
            },
        )*
    ) => {
        pub enum Bytecode {
            $($variant($variant),)*
        }

        $(
            pub struct $variant {
                $( pub $field: $ty, )*
            }

            impl From<$variant> for Bytecode {
                fn from(value: $variant) -> Bytecode {
                    Bytecode::$variant(value)
                }
            }
        )*

        macro_rules! dispatch_bytecode {
            ($obj:expr, $item:ident => $block:expr) => {
                match $obj {
                    $(Bytecode::$variant($item) => $block,)*
                }
            };
        }
        pub(crate) use dispatch_bytecode;
    }
}

Note, how the define_bytecode! macro defines the dispatch_bytecode! macro. That is, one macro defining another macro. That this code works came as a surprise to me, as I did not expect that Rust allows to nest macro definitions.

Using this macro, we can list the variants in a single place:

define_bytecode!(
    OuterSequenceStart{..},
    OuterSequenceEnd{..},
    OuterSequenceItem{..},
    OuterRecordStart{..},
    OuterRecordField{..}, 
    OuterRecordEnd{..},
    ProgramEnd{..},
    PushI32{..},
    PushF32{..},
);

and can finally write the accept function as:

impl Interpreter {
    fn accept(&mut self, event: &Event) -> Result<()> {
        dispatch_bytecode!(
            &self.program[self.current_instruction],
            instr => instr.accept(event),
        )
    }
}

See the code for the complete definition of the define_bytecode! macro, the bytecode used for serialization, or the bytecode used for deserialization.