dissect.cstruct 5.0.dev12__py3-none-any.whl → 5.0.dev14__py3-none-any.whl

dissect/cstruct/compiler.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import io
 import logging
+import struct
 from enum import Enum
 from itertools import groupby, repeat
 from keyword import iskeyword
@@ -11,7 +12,7 @@ from operator import itemgetter
 from textwrap import indent
 from typing import TYPE_CHECKING, Any
 
-from dissect.cstruct.exception import ExpressionParserError
+from dissect.cstruct.exception import ArraySizeError, ExpressionParserError
 from dissect.cstruct.expression import SOURCE_GLOBALS, Expression
 from dissect.cstruct.types import (
     Array,
@@ -30,7 +31,7 @@ from dissect.cstruct.types import (
     WcharArray,
 )
 from dissect.cstruct.types.enum import EnumMetaType
-from dissect.cstruct.types.structure import UnionMetaType, UnionProxy
+from dissect.cstruct.types.structure import StructureMetaType, UnionMetaType, UnionProxy
 from dissect.cstruct.util import ENDIANNESS_TO_BYTEORDER_MAP, _struct
 
 if TYPE_CHECKING:
@@ -63,25 +64,33 @@ python_compile = compile
 
 
 def compile_structure(cls: type[Structure]) -> type[Structure]:
-    """Compile the ``_read`` method of a structure or union type, if the compiler supports it.
+    """Compile the ``_read`` and ``_write`` methods of a structure or union type, if the compiler supports it.
 
     The compilation happens on first use.
     """
     try:
         _check_supported(cls, cls.__fields__)
         cls._read = _LazyCompileRead(cls)
+        cls._write = _LazyCompileWrite(cls)
         cls.__compiled__ = True
     except Exception as e:
-        # Revert to the interpreted reader
+        # Revert to the interpreted reader and writer
         log.debug("Failed to compile %s", cls, exc_info=e)
-        cls._read = classmethod(type(cls)._read)
-        cls.__compiled__ = False
+        _revert_to_interpreted(cls)
 
     return cls
 
 
-class _LazyCompileRead:
-    """Descriptor that compiles the ``_read`` method of a structure on first access and replaces itself with it."""
+def _revert_to_interpreted(cls: type[Structure]) -> None:
+    cls._read = classmethod(type(cls)._read)
+    cls._write = classmethod(type(cls)._write)
+    cls.__compiled__ = False
+
+
+class _LazyCompile:
+    """Descriptor that compiles a method of a structure on first access and replaces itself with it."""
+
+    name: str
 
     def __init__(self, cls: type[Structure]):
         self.cls = cls
@@ -103,18 +112,47 @@ class _LazyCompileRead:
             inline = all(type_.__dict__.get("__init__") is init for type_, init in self.snapshot.items())
 
             try:
-                self.method = _ReadSourceGenerator(cls, self.fields, inline).generate()
+                self.method = self._generator(inline).generate()
             except Exception as e:
-                # Revert to the slower loop based method
+                # Revert to the slower interpreted method
                 log.debug("Failed to compile %s", cls, exc_info=e)
-                self.method = classmethod(type(cls)._read)
-                cls.__compiled__ = False
+                self.method = self._revert()
 
-            if cls.__dict__.get("_read") is self:
-                cls._read = self.method
+            if cls.__dict__.get(self.name) is self:
+                setattr(cls, self.name, self.method)
 
         return self.method
 
+    def _generator(self, inline: bool) -> _SourceGenerator:
+        raise NotImplementedError
+
+    def _revert(self) -> classmethod:
+        """Return the interpreted version of the method."""
+        return classmethod(getattr(type(self.cls), self.name))
+
+
+class _LazyCompileRead(_LazyCompile):
+    """Descriptor that compiles the ``_read`` method of a structure on first access and replaces itself with it."""
+
+    name = "_read"
+
+    def _generator(self, inline: bool) -> _SourceGenerator:
+        return _ReadSourceGenerator(self.cls, self.fields, inline)
+
+    def _revert(self) -> classmethod:
+        # The structure no longer counts as compiled, so adding a field won't recompile _write either
+        _revert_to_interpreted(self.cls)
+        return super()._revert()
+
+
+class _LazyCompileWrite(_LazyCompile):
+    """Descriptor that compiles the ``_write`` method of a structure on first access and replaces itself with it."""
+
+    name = "_write"
+
+    def _generator(self, inline: bool) -> _SourceGenerator:
+        return _WriteSourceGenerator(self.cls, self.fields, inline)
+
 
 def _snapshot_structures(fields: list[Field], result: dict[type, Any]) -> dict[type, Any]:
     for field in fields:
@@ -166,7 +204,29 @@ _ENDIANNESS = ("<", ">", "!", "@", "=")
 # Types of which the value can be used as an integer directly
 _INTEGER_TYPES = (Int, Packed, Enum, Flag)
 
-# Globals that are available to all generated read functions
+# Types of which the value can be packed, unless they have a custom way of writing
+_PACKABLE_TYPES = (Array, Char, CharArray, Int, Packed, Pointer, Void, VoidArray, Wchar, WcharArray)
+
+
+def _field_bytes(type_: type[BaseType], value: Any, endian: str) -> bytes:
+    """Return the bytes of a char or wchar field that isn't of the right type yet or doesn't fit."""
+    buf = type_.dumps(type_.__default__() if value is None else value, endian=endian)
+    if len(buf) > type_.__size__:
+        raise ArraySizeError(f"Expected {type_.__size__} bytes for {type_.__name__}, got {len(buf)} instead.")
+    return buf
+
+
+def _array_values(type_: type[Array], value: list[Any] | None) -> list[Any]:
+    """Return the values of an array field that doesn't have the right size yet."""
+    return type_.__default__() if value is None else type_._pad(value)
+
+
+def _or_default(value: Any, type_: type[BaseType]) -> Any:
+    """Return the value of a field, or the default of its type if it's ``None``."""
+    return type_.__default__() if value is None else value
+
+
+# Globals that are available to all generated functions
 _GLOBALS = {
     "_type_call": type.__call__,
     "_object_new": object.__new__,
@@ -178,12 +238,20 @@ _GLOBALS = {
     "_getattr": getattr,
     "_UnionProxy": UnionProxy,
     "_union_read": UnionMetaType._read,
+    "_struct_write": StructureMetaType._write,
+    "_union_write": UnionMetaType._write,
+    "_struct_error": struct.error,
+    "_field_bytes": _field_bytes,
+    "_array_values": _array_values,
+    "_or_default": _or_default,
     **SOURCE_GLOBALS,
 }
 
 
-class _ReadSourceGenerator:
-    """Generates the source code of an optimized ``_read`` method for a structure or union."""
+class _SourceGenerator:
+    """Base class for generating the source code of an optimized method of a structure or union."""
+
+    name: str
 
     def __init__(self, cls: type[Structure], fields: list[Field], inline: bool = True):
         self.cls = cls
@@ -191,24 +259,17 @@ class _ReadSourceGenerator:
         self.align = cls.__align__
         self.inline = inline and not self.align
 
+        # The values available to the generated code
+        self.symbols: dict[str, Any] = {"cs": cls.__cs__}
+        self._symbol_names: dict[int, str] = {}
+
         # Witnesses of inlined structure types (symbol of the type -> symbol of its __init__ at compile time)
         self._witnesses: dict[str, str] = {}
-        self._fallback_symbol: str | None = None
 
         # The state of the bit buffer: the type of its current value and the remaining bits
         self._bits: tuple[type[BaseType], int] | None = None
         self._var_ids: dict[str, int] = {}
 
-        # The values available to the generated code
-        self.symbols: dict[str, Any] = {"cs": cls.__cs__}
-        self._symbol_names: dict[int, str] = {}
-
-        # (field name, local variable name) of all fields assigned so far
-        self._assigned: list[tuple[str, str]] = []
-        self._has_context = False
-        self._uses_start = False
-        self._uses_sizes = False
-
     def _type_symbol(self, type_: type[BaseType]) -> str:
         """Add a type as symbol, returns its name."""
         return self._add_symbol(type_, "type")
@@ -220,39 +281,12 @@ class _ReadSourceGenerator:
             self.symbols[symbol] = value
         return symbol
 
-    def _context(self, lines: list[str]) -> str:
-        """Return the name of the context dictionary, creating it from the fields read so far if needed."""
-        if not self._has_context:
-            self._has_context = True
-            items = ", ".join(f"{name!r}: {var}" for name, var in self._assigned)
-            lines.append(f"values = {{{items}}}")
-
-        return "values"
-
     def _var(self, name: str) -> str:
         """Return a unique variable name for the generated code."""
         index = self._var_ids.get(name, 0)
         self._var_ids[name] = index + 1
         return f"_{name}{index}"
 
-    def _assign(self, field: Field, value: str) -> str:
-        """Assign a value to a new local variable for the given field, and the context dictionary if it exists."""
-        var = f"_value{len(self._assigned)}"
-        self._assigned.append((field._name, var))
-        if self._has_context:
-            return f"{var} = values[{field._name!r}] = {value}"
-        return f"{var} = {value}"
-
-    def _assign_local(self, field: Field, value: str) -> list[str]:
-        return [self._assign(field, value)]
-
-    def _fallback(self) -> str:
-        """Return the symbol of the non-inlined fallback read function."""
-        if self._fallback_symbol is None:
-            fallback = _Fallback(self.cls, self.fields)
-            self._fallback_symbol = self._add_symbol(fallback, "fallback")
-        return self._fallback_symbol
-
     def _witness(self, type_: type[Structure]) -> str:
         """Add an inlined type as symbol, with a witness to detect modifications to it after compilation."""
         symbol = self._type_symbol(type_)
@@ -263,10 +297,22 @@ class _ReadSourceGenerator:
         """Return the conditions that are true if any of the inlined types were modified after compilation."""
         return [f"{type_}.__init__ is not {init}" for type_, init in self._witnesses.items()]
 
+    def _starts_bit_value(self, field: Field) -> bool:
+        """Return whether a bit field is the first of a new value of the bit buffer."""
+        return self._bits is None or self._bits[1] == 0 or self._bits[0] is not _get_bit_buffer_type(field)
+
+    def _advance_bits(self, field: Field) -> int:
+        """Update the state of the bit buffer for a bit field, returns the bits that were remaining before it."""
+        if self._starts_bit_value(field):
+            self._bits = (_get_bit_buffer_type(field), _get_bits_type(field).__size__ * 8)
+        remaining = self._bits[1]
+        self._bits = (self._bits[0], remaining - field.bits)
+        return remaining
+
     def generate(self) -> MethodType:
         source = self.generate_source()
 
-        code = python_compile(source, f"<compiled {self.cls.__name__}._read>", "exec")
+        code = python_compile(source, f"<compiled {self.cls.__name__}.{self.name}>", "exec")
         exec(code, {**_GLOBALS, **self.symbols}, d := {})
         obj = d.popitem()[1]
         obj.__source__ = source
@@ -278,6 +324,56 @@ class _ReadSourceGenerator:
             return self._generate_union_source()
         return self._generate_struct_source()
 
+    def _generate_struct_source(self) -> str:
+        raise NotImplementedError
+
+    def _generate_union_source(self) -> str:
+        raise NotImplementedError
+
+
+class _ReadSourceGenerator(_SourceGenerator):
+    """Generates the source code of an optimized ``_read`` method for a structure or union."""
+
+    name = "_read"
+
+    def __init__(self, cls: type[Structure], fields: list[Field], inline: bool = True):
+        super().__init__(cls, fields, inline)
+
+        self._fallback_symbol: str | None = None
+
+        # (field name, local variable name) of all fields assigned so far
+        self._assigned: list[tuple[str, str]] = []
+        self._has_context = False
+        self._uses_start = False
+        self._uses_sizes = False
+
+    def _context(self, lines: list[str]) -> str:
+        """Return the name of the context dictionary, creating it from the fields read so far if needed."""
+        if not self._has_context:
+            self._has_context = True
+            items = ", ".join(f"{name!r}: {var}" for name, var in self._assigned)
+            lines.append(f"values = {{{items}}}")
+
+        return "values"
+
+    def _assign(self, field: Field, value: str) -> str:
+        """Assign a value to a new local variable for the given field, and the context dictionary if it exists."""
+        var = f"_value{len(self._assigned)}"
+        self._assigned.append((field._name, var))
+        if self._has_context:
+            return f"{var} = values[{field._name!r}] = {value}"
+        return f"{var} = {value}"
+
+    def _assign_local(self, field: Field, value: str) -> list[str]:
+        return [self._assign(field, value)]
+
+    def _fallback(self) -> str:
+        """Return the symbol of the non-inlined fallback read function."""
+        if self._fallback_symbol is None:
+            fallback = _Fallback(self.cls, self.fields)
+            self._fallback_symbol = self._add_symbol(fallback, "fallback")
+        return self._fallback_symbol
+
     def _generate_struct_source(self) -> str:
         """Generate the code for reading a structure."""
         read_code = "\n".join(self._generate_fields_source())
@@ -405,6 +501,16 @@ class _ReadSourceGenerator:
 
             # Everything else - basic and composite types (and arrays of them), and inlined structures
             else:
+                if field.offset is None:
+                    # The alignment depends on the position in the stream, so this needs a new block
+                    if self.align:
+                        yield from flush()
+                # A block is read in order, so a field before the current position needs a new block
+                elif current_offset is None or field.offset < current_offset:
+                    yield from flush()
+                elif current_block:
+                    current_offset = field.offset
+
                 if not current_block:
                     yield from align_to_field(field)
                 current_block.append(field)
@@ -456,18 +562,11 @@ class _ReadSourceGenerator:
         element = self._type_symbol(field.type.__type__)
         return f"_type_call({type_}, {element}._read_array(stream, _count, context={context}, endian=endian))"
 
-    def _starts_bit_value(self, field: Field) -> bool:
-        """Return whether a bit field is read from a new value of the bit buffer."""
-        return self._bits is None or self._bits[1] == 0 or self._bits[0] is not _get_bits_type(field)
-
     def _generate_bits_source(self, field: Field) -> Iterator[str]:
         """Generate the code for reading a bit field."""
         read_type = _get_bits_type(field)
         new_value = self._starts_bit_value(field)
-        if new_value:
-            self._bits = (read_type, read_type.__size__ * 8)
-        remaining = self._bits[1]
-        self._bits = (read_type, remaining - field.bits)
+        remaining = self._advance_bits(field)
 
         lines = []
         if new_value:
@@ -498,7 +597,7 @@ class _ReadSourceGenerator:
         """Generate the code for reading consecutive packed fields."""
         unpack = _Unpack("data")
         reads = []
-        size = self._pack_fields(fields, 0, unpack, reads, self._assign_local, self.align)
+        size = self._pack_fields(fields, 0, unpack, reads, self._assign_local)
 
         lines = [f"buf = stream.read({size})", _check_eof(size)]
         if unpack.has_values():
@@ -513,7 +612,6 @@ class _ReadSourceGenerator:
         unpack: _Unpack,
         lines: list[str],
         target: Callable[[Field, str], list[str]],
-        align: bool = False,
     ) -> int:
         """Generate the code for reading the given fields from ``buf``, starting at the given offset.
 
@@ -523,7 +621,7 @@ class _ReadSourceGenerator:
         Returns the offset after the last field.
         """
         inline = _inline_kind if self.inline else None
-        for field, count, char in _generate_struct_info(fields, align, inline=inline):
+        for field, count, char in _generate_struct_info(fields, inline=inline):
             if field is None:
                 # Padding
                 unpack.add(count, "x")
@@ -595,16 +693,15 @@ class _ReadSourceGenerator:
         lines.append(f"_object_setattr({var}, '_endian', endian)")
         self._proxify(var, var, fields, lines)
 
-    def _proxify(self, union: str, var: str, fields: list[Field], lines: list[str]) -> None:
+    def _proxify(self, union: str, var: str, fields: list[Field], lines: list[str], attr: str | None = None) -> None:
         """Generate the code equivalent to ``Union._proxify``."""
         for field in fields:
             if issubclass(field.type, Structure):
                 nested = self._var("nested")
+                member = attr or field._name
                 lines.append(f"{nested} = _getattr({var}, {field._name!r})")
-                lines.append(
-                    f"_object_setattr({var}, {field._name!r}, _UnionProxy({union}, {field._name!r}, {nested}))"
-                )
-                self._proxify(union, nested, field.type.__fields__, lines)
+                lines.append(f"_object_setattr({var}, {field._name!r}, _UnionProxy({union}, {member!r}, {nested}))")
+                self._proxify(union, nested, field.type.__fields__, lines, member)
 
     def _generate_packed_field_source(
         self,
@@ -622,19 +719,18 @@ class _ReadSourceGenerator:
             array = self._type_symbol(field.type)
             element = self._type_symbol(field_type.__type__)
 
-            if issubclass(field_type.__type__, Int):
-                elem_size = field_type.__type__.__size__
-                signed = bool(field_type.__type__.signed)
+            if issubclass(read_type, Int):
+                # Everything that's read as an integer type is sliced directly out of the buffer
+                elem_size = read_type.__size__
+                signed = bool(read_type.signed)
                 lines.append(f"_array_buf = {getter}")
                 raw = f"_int_from_bytes(_array_buf[i:i + {elem_size}], _byteorder[endian], signed={signed})"
-                value = f"[_type_call({element}, {raw}) for i in range(0, {count}, {elem_size})]"
-            elif issubclass(field_type.__type__, Pointer):
+                getter = f"[{raw} for i in range(0, {count}, {elem_size})]"
+
+            if issubclass(field_type.__type__, Pointer):
                 context = self._context(lines)
                 item = f"{element}.__new__({element}, value, stream, context={context}, endian=endian)"
                 value = f"[{item} for value in {getter}]"
-            elif issubclass(read_type, (Wchar, Int)):
-                # Preserve the behaviour of the generic parsing path for these types
-                value = f"[{element}(value) for value in {getter}]"
             else:
                 value = f"list(map(_type_call, _repeat({element}), {getter}))"
 
@@ -647,7 +743,7 @@ class _ReadSourceGenerator:
             encoding = self._add_symbol(read_type.__encoding_map__, "encoding")
             # Using str() instead of .decode() also supports other bytes-like objects returned by the stream
             type_ = self._type_symbol(field.type)
-            lines.extend(target(field, f"_type_call({type_}, str({getter}, {encoding}[endian]))"))
+            lines.extend(target(field, f"_type_call({type_}, str({getter}, {encoding}[endian], 'surrogatepass'))"))
 
         elif issubclass(field_type, Pointer):
             context = self._context(lines)
@@ -694,8 +790,320 @@ class _Unpack:
 
     def functions(self, name: str) -> dict[str, Callable]:
         """Return the given function of the struct for every endianness."""
-        fmt = _optimize_struct_fmt(self.items)
-        return {endian: getattr(_struct(endian, fmt), name) for endian in _ENDIANNESS}
+        return _struct_functions(_optimize_struct_fmt(self.items), name)
+
+
+class _WriteSourceGenerator(_SourceGenerator):
+    """Generates the source code of an optimized ``_write`` method for a structure or union."""
+
+    name = "_write"
+
+    def __init__(self, cls: type[Structure], fields: list[Field], inline: bool = True):
+        super().__init__(cls, fields, inline)
+
+        self._uses_start = False
+
+    def _generate_method_source(self, interpreted: str, lines: Iterable[str]) -> str:
+        """Generate the code of the method, ``interpreted`` is the name of the method to fall back to."""
+        write_code = "\n".join(lines)
+
+        # Fall back to the generic method if any of the inlined types was modified after compilation
+        checks = [f"cls is not {self._add_symbol(self.cls, 'cls')}", *self._modified_checks()]
+
+        preamble = [f"if {' or '.join(checks)}: return {interpreted}(cls, stream, data, endian=endian)"]
+        if self._uses_start:
+            preamble.append("start = stream.tell()")
+
+        code = indent("\n".join([*preamble, write_code]), "    ")
+
+        return f"def _write(cls, stream, data, *, endian):\n{code}"
+
+    def _generate_struct_source(self) -> str:
+        """Generate the code for writing a structure."""
+        return self._generate_method_source("_struct_write", list(self._generate_fields_source()))
+
+    def _generate_union_source(self) -> str:
+        """Generate the code for writing a union."""
+        if not _can_pack(self.cls):
+            raise TypeError(f"Unsupported union for compiler: {self.cls}")
+
+        pack = _Pack()
+        self._pack_union("data", self.cls, pack)
+
+        lines = [self._generate_pack_source(pack), f"return {self.cls.__size__}"]
+        return self._generate_method_source("_union_write", lines)
+
+    def _generate_fields_source(self) -> Iterator[str]:
+        """Generate the code for writing all fields."""
+        # The offset of the current field in the structure, None if it's unknown due to a dynamic field
+        current_offset = 0
+        current_block: list[Field] = []
+        current_bits: list[Field] = []
+
+        def flush() -> Iterator[str]:
+            if current_block:
+                yield from self._generate_packed_source(current_block)
+                current_block[:] = []
+            if current_bits:
+                yield from self._generate_bits_source(current_bits)
+                current_bits[:] = []
+
+        def align_to_field(field: Field) -> Iterator[str]:
+            nonlocal current_offset
+
+            if field.offset is not None and field.offset != current_offset:
+                if current_offset is None:
+                    # This only writes padding if the stream is before the offset, otherwise the amount is negative
+                    yield f'stream.write(b"\\x00" * (start + {field.offset} - stream.tell()))'
+                else:
+                    yield f"stream.write({bytes(field.offset - current_offset)!r})"
+                    current_offset = field.offset
+
+            if self.align and field.offset is None:
+                yield f'stream.write(b"\\x00" * (-stream.tell() & ({field.alignment} - 1)))'
+
+        for field in self.fields:
+            if not field.bits:
+                self._bits = None
+                if current_bits:
+                    yield from flush()
+
+            # Unlike reading, we can't move back in the stream
+            if None not in (field.offset, current_offset) and field.offset < current_offset:
+                raise TypeError(f"Unsupported offset for compiler: {field}")
+
+            size = field.type.__size__
+
+            # Bit fields
+            if field.bits:
+                if self._starts_bit_value(field):
+                    if not _can_pack_bits(field):
+                        raise TypeError(f"Unsupported type for bit field: {field.type}")
+
+                    yield from flush()
+                    yield from align_to_field(field)
+                elif field.offset is not None:
+                    raise TypeError(f"Unsupported offset for bit field: {field}")
+                else:
+                    # Only the first bit field of a value moves the stream
+                    size = 0
+
+                self._advance_bits(field)
+                current_bits.append(field)
+
+            # Basic and composite types (and arrays of them), and inlined structures
+            elif _can_pack(field.type):
+                if field.offset is None:
+                    # The alignment depends on the position in the stream, so this needs a new block
+                    if self.align:
+                        yield from flush()
+                elif current_offset is None:
+                    yield from flush()
+                elif current_block:
+                    current_offset = field.offset
+
+                if not current_block:
+                    yield from align_to_field(field)
+                current_block.append(field)
+
+            # Everything else - sub structures, arrays of structures, multi-dimensional and dynamic arrays
+            else:
+                yield from flush()
+                yield from align_to_field(field)
+                yield from self._generate_field_source(field)
+
+                # An aligned structure writes more than its size if the stream isn't aligned
+                if self.align:
+                    size = None
+
+            if current_offset is not None:
+                current_offset = None if size is None else current_offset + size
+
+        yield from flush()
+
+        if self.align:
+            self._uses_start = True
+            yield 'stream.write(b"\\x00" * (-stream.tell() & (cls.__alignment__ - 1)))'
+
+        yield "return stream.tell() - start" if self._uses_start else f"return {current_offset}"
+
+    def _generate_field_source(self, field: Field) -> Iterator[str]:
+        """Generate the code for writing a field using its type."""
+        self._uses_start = True
+        type_ = self._type_symbol(field.type)
+        lines = [
+            f"_value = {_get_attribute('data', field._name)}",
+            f"if _value is None: _value = {type_}.__default__()",
+            f"{type_}._write(stream, _value, endian=endian)",
+        ]
+
+        yield _block(lines)
+
+    def _generate_bits_source(self, fields: list[Field]) -> Iterator[str]:
+        """Generate the code for writing the bit fields that share a value of the bit buffer."""
+        type_ = _get_bits_type(fields[0])
+
+        pack = _Pack()
+        little = []
+        other = []
+        remaining = type_.__size__ * 8
+        for field in fields:
+            value = self._load_value(field, _get_attribute("data", field._name), pack)
+            if issubclass(field.type, (Enum, Flag)):
+                value += "._value_"
+
+            remaining -= field.bits
+            little.append(f"{value} << {type_.__size__ * 8 - remaining - field.bits}")
+            other.append(f"{value} << {remaining}")
+
+        # The bit buffer only writes the least significant bits first for explicit little endian
+        self._pack_number(type_, f"({' | '.join(little)} if endian == '<' else {' | '.join(other)})", pack)
+
+        yield self._generate_pack_source(pack)
+
+    def _generate_packed_source(self, fields: list[Field]) -> Iterator[str]:
+        """Generate the code for writing consecutive packed fields."""
+        pack = _Pack()
+        self._pack_fields(fields, "data", pack)
+
+        if pack.items:
+            yield self._generate_pack_source(pack)
+
+    def _generate_pack_source(self, pack: _Pack) -> str:
+        """Generate the code for checking, packing and writing values."""
+        function = f"{self._add_symbol(pack.functions('pack'), 'pack')}[endian]"
+        if pack.values == pack.retry:
+            return _block([*pack.lines, f"stream.write({function}({', '.join(pack.values)}))"])
+
+        # Numbers aren't checked for None up front, since struct.pack raises an error for those
+        return _block(
+            [
+                *pack.lines,
+                "try:",
+                f"    buf = {function}({', '.join(pack.values)})",
+                "except _struct_error:",
+                f"    buf = {function}({', '.join(pack.retry)})",
+                "stream.write(buf)",
+            ]
+        )
+
+    def _load_value(self, field: Field, value: str, pack: _Pack) -> str:
+        """Generate the code for loading a value, using the default of its type if it's ``None``."""
+        var = self._var("value")
+        pack.lines.append(f"if ({var} := {value}) is None: {var} = {self._type_symbol(field.type)}.__default__()")
+        return var
+
+    def _pack_fields(self, fields: list[Field], var: str, pack: _Pack, proxied: bool = False) -> None:
+        """Generate the code for packing the given fields of the object in ``var``.
+
+        The fields of objects in a union are proxied, in which case we use the object behind the proxy.
+        """
+        for field, count, char in _generate_struct_info(fields, inline=_inline_kind):
+            if field is None:
+                # Padding
+                pack.add(count, "x")
+                continue
+
+            value = _get_attribute(var, field._name)
+
+            if char in ("S", "U"):
+                # Inlined structure or union
+                self._witness(field.type)
+                nested = self._load_value(field, value, pack)
+                if proxied:
+                    pack.lines.append(f"if type({nested}) is _UnionProxy: {nested} = {nested}.__target__")
+
+                if char == "S":
+                    self._pack_fields(field.type.__fields__, nested, pack, proxied)
+                else:
+                    self._pack_union(nested, field.type, pack)
+                continue
+
+            self._pack_field(field, value, count, char, pack)
+
+    def _pack_union(self, var: str, type_: type[Union], pack: _Pack) -> None:
+        """Generate the code for packing the union in ``var``, by the field the interpreted writer uses."""
+        # Changing a nested structure can change which field is the largest
+        for field in type_.__fields__:
+            if issubclass(field.type, Structure):
+                self._witness(field.type)
+
+        self._pack_fields([type_._write_field()], var, pack, proxied=True)
+
+    def _pack_field(self, field: Field, value: str, count: int, char: str, pack: _Pack) -> None:
+        """Generate the code for checking and packing a field."""
+        field_type = field.type
+        read_type = _get_read_type(field_type)
+        type_ = self._type_symbol(field_type)
+
+        if issubclass(field_type, Array):
+            read_type = _get_read_type(field_type.__type__)
+            var = self._var("value")
+            num = field_type.__num_entries__
+            pack.lines.append(
+                f"if ({var} := {value}) is None or len({var}) != {num}: {var} = _array_values({type_}, {var})"
+            )
+
+            # Entries can be enum members or plain integers, both can be packed as they are
+            if issubclass(read_type, Int):
+                pack.add(count, "s", f"b''.join([{_int_to_bytes(read_type, 'item')} for item in {var}])")
+            else:
+                pack.add(count, char, f"*{var}")
+
+        elif issubclass(field_type, (Char, CharArray, Wchar, WcharArray)):
+            var = self._var("value")
+
+            # The format pads values that are too short with zeroes
+            if issubclass(field_type, (Wchar, WcharArray)):
+                encoding = self._add_symbol(Wchar.__encoding_map__, "encoding")
+                pack.lines.append(
+                    f"if ({var} := {value}) is None or "
+                    f"len({var} := {var}.encode({encoding}[endian], 'surrogatepass')) > {count}: "
+                    f"{var} = _field_bytes({type_}, {value}, endian)"
+                )
+            else:
+                pack.lines.append(
+                    f"if not isinstance({var} := {value}, bytes) or len({var}) > {count}: "
+                    f"{var} = _field_bytes({type_}, {var}, endian)"
+                )
+
+            pack.add(count, "s", var)
+
+        elif issubclass(field_type, (Enum, Flag)):
+            self._pack_number(read_type, f"{self._load_value(field, value, pack)}._value_", pack)
+
+        elif issubclass(read_type, Packed):
+            pack.add(1, read_type.__fmt__, value, f"_or_default({value}, {type_})")
+
+        else:
+            self._pack_number(read_type, self._load_value(field, value, pack), pack)
+
+    def _pack_number(self, type_: type[BaseType], value: str, pack: _Pack) -> None:
+        """Generate the code for packing a value of an integer or packed type."""
+        if issubclass(type_, Int):
+            pack.add(type_.__size__, "s", _int_to_bytes(type_, value))
+        else:
+            pack.add(1, type_.__fmt__, value)
+
+
+class _Pack:
+    """Tracks consecutive struct pack formats, and the code for the values to pack."""
+
+    def __init__(self):
+        self.items: list[tuple[int, str]] = []
+        self.lines: list[str] = []
+        self.values: list[str] = []
+        self.retry: list[str] = []
+
+    def add(self, count: int, char: str, value: str | None = None, retry: str | None = None) -> None:
+        self.items.append((count, char))
+        if value is not None:
+            self.values.append(value)
+            self.retry.append(retry or value)
+
+    def functions(self, name: str) -> dict[str, Callable]:
+        """Return the given function of the struct for every endianness."""
+        return _struct_functions(_optimize_struct_fmt(self.items), name)
 
 
 def _check_union(cls: type[Union], fields: list[Field], inline: bool = True) -> int:
@@ -767,6 +1175,52 @@ def _uses_generated_init(type_: type[Structure]) -> bool:
     return "__init__" in type_.__dict__ and type_.__new__ is object.__new__
 
 
+def _can_pack(type_: type[BaseType]) -> bool:
+    """Determine if a value of a type can be written by packing it, instead of by using the ``_write`` method."""
+    if type_.__size__ is None:
+        return False
+
+    if issubclass(type_, (Enum, Flag)):
+        return _can_pack(type_.__type__)
+
+    if issubclass(type_, Structure):
+        kind = _inline_kind(type_)
+        if kind == "U":
+            field = type_._write_field()
+            return field is not None and bool(field.type.__size__) and _can_pack(field.type)
+        return kind == "S" and all(_can_pack(field.type) for field in type_.__fields__)
+
+    if issubclass(type_, Array) and (
+        issubclass(type_.__type__, (Structure, BaseArray)) or not _can_pack(type_.__type__)
+    ):
+        return False
+
+    return any(issubclass(type_, base) and _has_default_write(type_, base) for base in _PACKABLE_TYPES)
+
+
+def _can_pack_bits(field: Field) -> bool:
+    """Determine if the value of the bit buffer for a bit field can be written by packing it."""
+    type_ = _get_bits_type(field)
+    return issubclass(type_, (Int, Packed)) and _can_pack(type_)
+
+
+def _has_default_write(type_: type[BaseType], base: type[BaseType]) -> bool:
+    """Return whether a type uses the write methods of its base type, and doesn't have custom ones."""
+    return all(getattr(type_, name).__func__ is getattr(base, name).__func__ for name in ("_write", "_write_array"))
+
+
+def _int_to_bytes(type_: type[Int], value: str) -> str:
+    """Return the code to convert a value of an integer type to bytes."""
+    return f"{value}.to_bytes({type_.__size__}, _byteorder[endian], signed={bool(type_.signed)})"
+
+
+def _get_attribute(var: str, name: str) -> str:
+    """Return the code to get an attribute."""
+    if name.isidentifier() and not iskeyword(name):
+        return f"{var}.{name}"
+    return f"_getattr({var}, {name!r})"
+
+
 def _set_attribute(var: str, name: str, value: str) -> str:
     """Return the code to set an attribute."""
     if name.isidentifier() and not iskeyword(name):
@@ -809,25 +1263,18 @@ def _block(lines: list[str]) -> str:
 
 def _generate_struct_info(
     fields: list[Field],
-    align: bool = False,
     inline: Callable[[type[BaseType]], str | None] | None = None,
 ) -> Iterator[tuple[Field, int, str]]:
     if not fields:
         return
 
     current_offset = fields[0].offset
-    imaginary_offset = 0
     for field in fields:
         # We moved -- probably due to alignment
         if field.offset is not None and (drift := field.offset - current_offset) > 0:
             yield None, drift, "x"
             current_offset += drift
 
-        if align and field.offset is None and (drift := -imaginary_offset & (field.alignment - 1)) > 0:
-            # Assume we started at a correctly aligned boundary
-            yield None, drift, "x"
-            imaginary_offset += drift
-
         count = 1
         read_type = _get_read_type(field.type)
 
@@ -853,25 +1300,41 @@ def _generate_struct_info(
         elif issubclass(read_type, (Char, Wchar, Int)):
             yield field, count * read_type.__size__, "x"
 
-        size = count * read_type.__size__
-        imaginary_offset += size
         if current_offset is not None:
-            current_offset += size
+            current_offset += count * read_type.__size__
 
 
 def _optimize_struct_fmt(items: Iterable[tuple[int, str]]) -> str:
-    """Create a struct format string from ``(count, char)`` items, merging consecutive items of the same char."""
+    """Create a struct format string from ``(count, char)`` items, merging consecutive items of the same char.
+
+    Byte strings aren't merged, since those are a single value regardless of their count.
+    """
     chars = []
     for char, group in groupby(items, key=itemgetter(1)):
-        if count := sum(count for count, _ in group):
+        if char == "s":
+            chars.extend(f"{count}{char}" for count, _ in group)
+        elif count := sum(count for count, _ in group):
             chars.append(f"{count if count > 1 else ''}{char}")
 
     return "".join(chars)
 
 
+def _struct_functions(fmt: str, name: str) -> dict[str, Callable]:
+    """Return the given function of a struct with the given format for every endianness.
+
+    The ``@`` endianness uses the ``=`` format, since ``@`` would add padding between items to align them.
+    """
+    return {endian: getattr(_struct("=" if endian == "@" else endian, fmt), name) for endian in _ENDIANNESS}
+
+
+def _get_bit_buffer_type(field: Field) -> type[BaseType]:
+    """Return the type that ``BitBuffer`` uses for a bit field, only bit fields of the same type share a value."""
+    return field.type.__type__ if issubclass(field.type, (Enum, Flag)) else field.type
+
+
 def _get_bits_type(field: Field) -> type[BaseType]:
     """Return the type that ``BitBuffer`` reads a bit field as."""
-    type_ = field.type.__type__ if issubclass(field.type, (Enum, Flag)) else field.type
+    type_ = _get_bit_buffer_type(field)
     return type_.__cs__.uint8 if issubclass(type_, Char) else type_
 
 

dissect/cstruct/cstruct.py

@@ -86,7 +86,7 @@ class cstruct:
             "int24": self._make_int_type("int24", 3, True, alignment=4),
             "uint24": self._make_int_type("uint24", 3, False, alignment=4),
             "int48": self._make_int_type("int48", 6, True, alignment=8),
-            "uint48": self._make_int_type("int48", 6, False, alignment=8),
+            "uint48": self._make_int_type("uint48", 6, False, alignment=8),
             "int128": self._make_int_type("int128", 16, True, alignment=16),
             "uint128": self._make_int_type("uint128", 16, False, alignment=16),
 

dissect/cstruct/types/base.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import functools
 import warnings
 from io import SEEK_CUR, BytesIO
+from types import MethodType
 from typing import TYPE_CHECKING, Any, BinaryIO, ClassVar, TypeVar
 
 from dissect.cstruct.exception import ArraySizeError
@@ -19,6 +20,8 @@ if TYPE_CHECKING:
 
 EOF = -0xE0F  # Negative counts are illegal anyway, so abuse that for our EOF sentinel
 
+_IMMUTABLE_TYPES = (int, float, bytes, str)
+
 
 class MetaType(type):
     """Base metaclass for cstruct type classes."""
@@ -218,7 +221,8 @@ class MetaType(type):
             array: The array to write.
             endian: The endianness to use when parsing (little, big, network, <, >, !, @ or =).
         """
-        return sum(cls._write(stream, entry, endian=endian) for entry in array)
+        write = cls._write
+        return sum(write(stream, entry, endian=endian) for entry in array)
 
     def _write_0(cls, stream: BinaryIO, array: list[Self], *, endian: Endianness) -> int:  # type: ignore
         """Internal function for writing null-terminated arrays.
@@ -295,20 +299,29 @@ class _overload:
         b'\\x7b\\x00\\x00\\x00'
     """
 
-    def __init__(self, func: Callable[..., Any]) -> None:
+    def __init__(self, func: Callable[..., Any], method: Callable[..., Any]) -> None:
         self.func = func
+        self.method = method
 
-    def __get__(self, instance: BaseType | None, owner: type[BaseType]) -> Callable[[], bytes]:
+    def __get__(self, instance: BaseType | None, owner: type[BaseType]) -> Callable[..., Any]:
         if instance is None:
             return functools.partial(self.func, owner)
-        return functools.partial(self.func, instance.__class__, value=instance)
+        return MethodType(self.method, instance)
+
+
+def _instance_dumps(self: BaseType, *, endian: AllowedEndianness | None = None) -> bytes:
+    return MetaType.dumps(self.__class__, self, endian=endian)
+
+
+def _instance_write(self: BaseType, stream: BinaryIO, *, endian: AllowedEndianness | None = None) -> int:
+    return MetaType.write(self.__class__, stream, self, endian=endian)
 
 
 class BaseType(metaclass=MetaType):
     """Base class for cstruct type classes."""
 
-    dumps = _overload(MetaType.dumps)
-    write = _overload(MetaType.write)
+    dumps = _overload(MetaType.dumps, _instance_dumps)
+    write = _overload(MetaType.write, _instance_write)
 
     def __len__(self) -> int:
         """Return the byte size of the type."""
@@ -338,9 +351,10 @@ class BaseArray(BaseType):
 
     @classmethod
     def __default__(cls) -> BaseType:
-        return type.__call__(
-            cls, [cls.__type__.__default__()] * (cls.__num_entries__ if isinstance(cls.__num_entries__, int) else 0)
-        )
+        count = cls.__num_entries__ if isinstance(cls.__num_entries__, int) else 0
+        if isinstance(value := cls.__type__.__default__(), _IMMUTABLE_TYPES):
+            return type.__call__(cls, [value] * count)
+        return type.__call__(cls, [cls.__type__.__default__() for _ in range(count)])
 
     @classmethod
     def _read(cls, stream: BinaryIO, *, context: dict[str, Any] | None = None, endian: Endianness) -> list[BaseType]:
@@ -370,12 +384,24 @@ class BaseArray(BaseType):
         if cls.__null_terminated__:
             return cls.__type__._write_0(stream, data, endian=endian)
 
-        if not cls.__dynamic__ and (remaining := cls.__num_entries__ - (actual_size := len(data))):
+        if not cls.__dynamic__:
+            data = cls._pad(data)
+
+        return cls.__type__._write_array(stream, data, endian=endian)
+
+    @classmethod
+    def _pad(cls, data: list[Any]) -> list[Any]:
+        """Internal function for padding the entries of a static array to its size.
+
+        Args:
+            data: The entries of the array.
+        """
+        if remaining := cls.__num_entries__ - (actual_size := len(data)):
             if remaining < 0:
                 raise ArraySizeError(f"Expected static array size {cls.__num_entries__}, got {actual_size} instead.")
             data = data + [cls.__type__.__default__()] * remaining
 
-        return cls.__type__._write_array(stream, data, endian=endian)
+        return data
 
     @classproperty
     def type(cls) -> type[BaseType]:

dissect/cstruct/types/char.py

@@ -23,8 +23,9 @@ class CharArray(bytes, BaseArray):
         return type.__call__(cls, super()._read(stream, context=context, endian=endian))
 
     @classmethod
-    def _write(cls, stream: BinaryIO, data: bytes | str | list[int], *, endian: Endianness) -> int:
-        if isinstance(data, list):
+    def _write(cls, stream: BinaryIO, data: bytes | bytearray | str | list[int], *, endian: Endianness) -> int:
+        # Copy mutable data types to bytes so we don't modify them in place
+        if isinstance(data, (list, bytearray)):
             data = bytes(data)
         elif isinstance(data, str):
             data = data.encode("latin-1")
@@ -92,4 +93,7 @@ class Char(bytes, BaseType):
         if isinstance(data, str):
             data = data.encode("latin-1")
 
-        return stream.write(data)
+        if (actual_size := len(data)) > 1:
+            raise ValueError(f"Expected a single character, got {actual_size} instead.")
+
+        return stream.write(data or b"\x00")

dissect/cstruct/types/structure.py

@@ -7,7 +7,7 @@ from collections.abc import MutableMapping
 from contextlib import contextmanager
 from enum import Enum
 from functools import lru_cache
-from itertools import chain
+from itertools import chain, repeat
 from operator import attrgetter
 from textwrap import dedent
 from types import FunctionType
@@ -18,6 +18,7 @@ from dissect.cstruct.expression import Expression
 from dissect.cstruct.lexer import IDENTIFIER_TYPES
 from dissect.cstruct.types.base import (
     _BUFFER_TYPES,
+    _IMMUTABLE_TYPES,
     BaseArray,
     BaseType,
     MetaType,
@@ -141,7 +142,6 @@ class StructureMetaType(MetaType):
         else:
             size, alignment = cls._calculate_size_and_offsets(fields, align)
 
-        # TODO: compile _write
         # TODO: generate cached_property for lazy reading
 
         classdict["__size__"] = size
@@ -193,7 +193,7 @@ class StructureMetaType(MetaType):
                     # Moved to a bit field of another type, e.g. uint16 f1 : 8, uint32 f2 : 8;
                     or field_type != bits_type
                     # Still processing a bit field, but it's at a different offset due to alignment or a manual offset
-                    or (bits_type is not None and offset > bits_field_offset + bits_type.__size__)
+                    or (None not in (offset, bits_field_offset) and offset > bits_field_offset + bits_type.__size__)
                 ):
                     # ... if any of this is true, we have to move to the next field
                     bits_type = field_type
@@ -299,7 +299,6 @@ class StructureMetaType(MetaType):
     def _write(cls, stream: BinaryIO, data: Structure, *, endian: Endianness) -> int:
         bit_buffer = None
         struct_start = stream.tell()
-        num = 0
 
         for field in cls.__fields__:
             field_type = field.type
@@ -331,7 +330,7 @@ class StructureMetaType(MetaType):
                     stream.write(b"\x00" * align_pad)
                     offset += align_pad
 
-            value = getattr(data, field._name, None)
+            value = getattr(data, field._name)
             if value is None:
                 value = field_type.__default__()
 
@@ -345,7 +344,6 @@ class StructureMetaType(MetaType):
                     bit_buffer.write(field_type, value, field.bits)
             else:
                 field_type._write(stream, value, endian=endian)
-                num += stream.tell() - offset
 
         if bit_buffer is not None and bit_buffer._type is not None:
             bit_buffer.flush()
@@ -354,7 +352,7 @@ class StructureMetaType(MetaType):
             # Align the stream
             stream.write(b"\x00" * (-stream.tell() & (cls.__alignment__ - 1)))
 
-        return num
+        return stream.tell() - struct_start
 
     def add_field(cls, name: str, type_: type[BaseType], bits: int | None = None, offset: int | None = None) -> None:
         field = Field(name, type_, bits=bits, offset=offset)
@@ -698,6 +696,10 @@ class UnionMetaType(StructureMetaType):
 
         return obj
 
+    def _write_field(cls) -> Field | None:
+        """Return the field to write a union by, its largest field and preferably not an anonymous structure."""
+        return max(cls.__fields__, key=lambda e: (e.type.__size__ or 0, e.name is not None), default=None)
+
     def _write(cls, stream: BinaryIO, data: Union, *, endian: Endianness) -> int:
         if cls.__dynamic__:
             raise NotImplementedError("Writing dynamic unions is not yet supported")
@@ -705,24 +707,9 @@ class UnionMetaType(StructureMetaType):
         offset = stream.tell()
         expected_offset = offset + len(cls)
 
-        # Sort by largest field
-        fields = sorted(cls.__fields__, key=lambda e: e.type.__size__ or 0, reverse=True)
-        anonymous_struct = False
-
-        # Try to write by largest field
-        for field in fields:
-            if isinstance(field.type, StructureMetaType) and field.name is None:
-                # Prefer to write regular fields initially
-                anonymous_struct = field.type
-                continue
-
+        if (field := cls._write_field()) is not None:
             # Write the value
             field.type._write(stream, getattr(data, field._name), endian=endian)
-            break
-
-        # If we haven't written anything yet and we initially skipped an anonymous struct, write it now
-        if stream.tell() == offset and anonymous_struct:
-            anonymous_struct._write(stream, data, endian=endian)
 
         # If we haven't filled the union size yet, pad it
         if remaining := expected_offset - stream.tell():
@@ -745,7 +732,9 @@ class Union(Structure, metaclass=UnionMetaType):
             raise NotImplementedError("Modifying a dynamic union is not yet supported")
 
         super().__setattr__(attr, value)
-        self._rebuild(attr)
+        # Fields of an anonymous structure are rebuilt through the proxy of that structure
+        if attr in self.__class__.__lookup__:
+            self._rebuild(attr)
 
     def _rebuild(self, attr: str) -> None:
         if (cur_buf := getattr(self, "_buf", None)) is None:
@@ -773,13 +762,17 @@ class Union(Structure, metaclass=UnionMetaType):
         object.__setattr__(self, "__dynamic_sizes__", sizes)
 
     def _proxify(self) -> None:
-        def _proxy_structure(value: Structure) -> None:
+        def _proxy_structure(value: Structure, attr: str | None = None) -> None:
             for field in value.__class__.__fields__:
                 if issubclass(field.type, Structure):
                     nested_value = getattr(value, field._name)
-                    proxy = UnionProxy(self, field._name, nested_value)
-                    object.__setattr__(value, field._name, proxy)
-                    _proxy_structure(nested_value)
+                    # Changing a nested field rebuilds the union from the top level field it's in
+                    member = attr or field._name
+                    object.__setattr__(value, field._name, UnionProxy(self, member, nested_value))
+                    # The value of a nested union is already a proxy, so continue with the value behind it
+                    if isinstance(nested_value, UnionProxy):
+                        nested_value = nested_value.__target__
+                    _proxy_structure(nested_value, member)
 
         _proxy_structure(self)
 
@@ -864,7 +857,7 @@ def _make_structure__init__(fields: list[str]) -> str:
     """
     field_args = ", ".join(f"{field} = None" for field in fields)
     field_init = "\n".join(
-        f" self.{name} = {name} if {name} is not None else _{i}_default" for i, name in enumerate(fields)
+        f" self.{name} = {name} if {name} is not None else _{i}_default()" for i, name in enumerate(fields)
     )
 
     code = f"def __init__(self{', ' + field_args or ''}):\n"
@@ -880,7 +873,7 @@ def _make_union__init__(fields: list[str]) -> str:
     """
     field_args = ", ".join(f"{field} = None" for field in fields)
     field_init = "\n".join(
-        f" object.__setattr__(self, '{name}', {name} if {name} is not None else _{i}_default)"
+        f" object.__setattr__(self, '{name}', {name} if {name} is not None else _{i}_default())"
         for i, name in enumerate(fields)
     )
 
@@ -976,7 +969,7 @@ def _generate_structure__init__(fields: list[Field]) -> FunctionType:
             co_names=_remap_co_values(template.__code__.co_names, mapping),
             co_varnames=_remap_co_values(template.__code__.co_varnames, mapping),
         ),
-        template.__globals__ | {f"__{field._name}_default__": field.type.__default__() for field in fields},
+        template.__globals__ | {f"__{field._name}_default__": _default_factory(field.type) for field in fields},
         argdefs=template.__defaults__,
     )
 
@@ -996,11 +989,17 @@ def _generate_union__init__(fields: list[Field]) -> FunctionType:
             co_names=_remap_co_values(template.__code__.co_names, mapping),
             co_varnames=_remap_co_values(template.__code__.co_varnames, mapping),
         ),
-        template.__globals__ | {f"__{field._name}_default__": field.type.__default__() for field in fields},
+        template.__globals__ | {f"__{field._name}_default__": _default_factory(field.type) for field in fields},
         argdefs=template.__defaults__,
     )
 
 
+def _default_factory(type_: type[BaseType]) -> Callable[[], Any]:
+    """Return a function that returns default values of a type, which are only shared if they're immutable."""
+    value = type_.__default__()
+    return repeat(value).__next__ if isinstance(value, _IMMUTABLE_TYPES) else type_.__default__
+
+
 def _generate_co_mapping(fields: list[Field]) -> dict[str, str]:
     """Generates a mapping of generated code object names to field names.
 

dissect/cstruct/types/wchar.py

@@ -28,12 +28,15 @@ class WcharArray(str, BaseArray):
         if cls.__null_terminated__:
             data += "\x00"
 
-        if not cls.__dynamic__ and (remaining := cls.__num_entries__ - (actual_size := len(data))):
+        buf = data.encode(Wchar.__encoding_map__[endian], "surrogatepass")
+
+        # A surrogate pair is a single character that takes up two entries of the array
+        if not cls.__dynamic__ and (remaining := cls.__num_entries__ - (actual_size := len(buf) // 2)):
             if remaining < 0:
                 raise ArraySizeError(f"Expected static array size {cls.__num_entries__}, got {actual_size} instead.")
-            data += "\x00" * remaining
+            buf += b"\x00\x00" * remaining
 
-        return stream.write(data.encode(Wchar.__encoding_map__[endian]))
+        return stream.write(buf)
 
 
 class Wchar(str, BaseType):
@@ -42,6 +45,7 @@ class Wchar(str, BaseType):
     __ArrayType__ = WcharArray
 
     __slots__ = ()
+    # Decode with "surrogatepass" so lone surrogates are handled correctly
     __encoding_map__: ClassVar[dict[Endianness, str]] = {
         "@": f"utf-16-{sys.byteorder[0]}e",
         "=": f"utf-16-{sys.byteorder[0]}e",
@@ -72,12 +76,13 @@ class Wchar(str, BaseType):
         if count != EOF and len(data) != count:
             raise EOFError(f"Read {len(data)} bytes, but expected {count}")
 
-        return type.__call__(cls, str(data, cls.__encoding_map__[endian]))
+        return type.__call__(cls, str(data, cls.__encoding_map__[endian], "surrogatepass"))
 
     @classmethod
     def _read_0(cls, stream: BinaryIO, *, context: dict[str, Any] | None = None, endian: Endianness) -> Wchar:
         if _is_seekable(stream):
-            return type.__call__(cls, _read_null_terminated(stream, 2).decode(cls.__encoding_map__[endian]))
+            data = _read_null_terminated(stream, 2)
+            return type.__call__(cls, data.decode(cls.__encoding_map__[endian], "surrogatepass"))
 
         buf = []
         while True:
@@ -90,8 +95,14 @@ class Wchar(str, BaseType):
 
             buf.append(point)
 
-        return type.__call__(cls, b"".join(buf).decode(cls.__encoding_map__[endian]))
+        return type.__call__(cls, b"".join(buf).decode(cls.__encoding_map__[endian], "surrogatepass"))
 
     @classmethod
     def _write(cls, stream: BinaryIO, data: str, *, endian: Endianness) -> int:
-        return stream.write(data.encode(cls.__encoding_map__[endian]))
+        buf = data.encode(cls.__encoding_map__[endian], "surrogatepass")
+
+        # A surrogate pair is a single character that doesn't fit in a single wchar
+        if (actual_size := len(buf) // 2) > 1:
+            raise ValueError(f"Expected a single character, got {actual_size} instead.")
+
+        return stream.write(buf or b"\x00\x00")

{dissect_cstruct-5.0.dev12.dist-info → dissect_cstruct-5.0.dev14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dissect.cstruct
-Version: 5.0.dev12
+Version: 5.0.dev14
 Summary: A Dissect module implementing a parser for C-like structures: structure parsing in Python made easy
 Author-email: Dissect Team <dissect@fox-it.com>
 License-Expression: Apache-2.0

{dissect_cstruct-5.0.dev12.dist-info → dissect_cstruct-5.0.dev14.dist-info}/RECORD

@@ -1,7 +1,7 @@
 dissect/cstruct/__init__.py,sha256=KOyNeUvu4c0JaAiyOfj7iftGwUIjTmD8MwHtFsLPLSU,1504
 dissect/cstruct/bitbuffer.py,sha256=aNzY4lkWe8ffAc8j7FODXzwZc8gZwbey77UosyoYCsw,2711
-dissect/cstruct/compiler.py,sha256=hhX1bxZFsPOcs6BUO82YHNtWqhKsiOb-H-d18Aa1Q6E,35148
-dissect/cstruct/cstruct.py,sha256=VQDgGF7OUaPKQ_Wsy_2fUHL9Jlims6R2JOh-1_bgiZ8,26550
+dissect/cstruct/compiler.py,sha256=0u58J5pB4izosBj8G4wiBI1JSAAgHFMLM63IaSG4YO4,54509
+dissect/cstruct/cstruct.py,sha256=mJerzPouWV29RdK5LqjGitDDMvNKcadoT2hj1zpExnM,26551
 dissect/cstruct/exception.py,sha256=RqD08LuBkbMqa3iQ1A05DNN8XgEmQR_SUnguwe9LGo8,316
 dissect/cstruct/expression.py,sha256=Y27-N48BomBoRFozbWhGDF_DS2lfklRwWVUbCzGnxHE,10873
 dissect/cstruct/lexer.py,sha256=NDUfAWgyxgRNRuVO8ExXQ51LqQXS2mv9bywVa_M7tFQ,19665
@@ -10,21 +10,21 @@ dissect/cstruct/util.py,sha256=vAf-kdOkk3zprW_gD6QmQeN3VIRgQqhPS4B_ha1jEpQ,18040
 dissect/cstruct/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dissect/cstruct/tools/stubgen.py,sha256=MeZkLmSeNSLNcAEwlAdIGv1lKvlvK3UQ4WN85ovsxTM,7672
 dissect/cstruct/types/__init__.py,sha256=6P71Pu94Dv7RNjZpDZ19pqzx-G0QzmQw4q0-3zWALso,911
-dissect/cstruct/types/base.py,sha256=10HLF2SDeWUMVc6ISMv8xKdEggqrxLPHnJMwI03DEFg,17216
-dissect/cstruct/types/char.py,sha256=Rzjn2-pXbn-fV_UnNRGgQugFZ50mPuXQhcmgHvY5PR0,3172
+dissect/cstruct/types/base.py,sha256=bbaxiHCVkoU6ugqn5RYishPQ5L6xGrPfQuejjro0K_g,18133
+dissect/cstruct/types/char.py,sha256=N19R7Zzn2uMuG-ULW5X2UgmgkLvjaBj1Td7AQQI93Fg,3417
 dissect/cstruct/types/enum.py,sha256=nOdI6vGtqTfPMhOBkKuiuOriIB2YHSt8Q9ruri37s3U,9177
 dissect/cstruct/types/flag.py,sha256=YxQHix8ZhK1qQX3J_Og7GJp_4JAxCwzUbUibuJeDaWQ,3479
 dissect/cstruct/types/int.py,sha256=9RbdrkvAI28zATM97VVORQnras0XhLaVMtdPjJwpTH4,2089
 dissect/cstruct/types/leb128.py,sha256=V97XTJT_Xbw1LpTwTN1OkjPHuQQAQUxd8OGuA6yEDmE,2387
 dissect/cstruct/types/packed.py,sha256=NgCX0YWX8SH8z47atq0f39w9MBjU_eBFJtUisUWdO-M,3867
 dissect/cstruct/types/pointer.py,sha256=FNuB9bfENWRebbKX-R6Vs6xWPMHepjvOhhomeUhr7I8,3985
-dissect/cstruct/types/structure.py,sha256=WDMmfhP2XE5p1Bn-WgFFstfaWC709AU7_y-SGi3TyDw,39116
+dissect/cstruct/types/structure.py,sha256=mxDvurQ0oAjI2flykhAPlT-03280cDpwGTz0-bqbr-0,39558
 dissect/cstruct/types/void.py,sha256=fV4hKk50uop_EKbR1nfy1cdtAzvUYrSymXsdakCyFk8,1222
-dissect/cstruct/types/wchar.py,sha256=qUbX6auj0qHWSOkgvaI7ZW_f7RGgwa-Hf_d_mDpntlM,3410
-dissect_cstruct-5.0.dev12.dist-info/licenses/COPYRIGHT,sha256=H-18RXfshdH9AdHwW2eO1Xa-5s6tY5eipHh5c0whDu4,316
-dissect_cstruct-5.0.dev12.dist-info/licenses/LICENSE,sha256=PhUqiw6jAh2KbBdVRPBq_hfAvfcTBin7nZ3CK7NQbTM,11341
-dissect_cstruct-5.0.dev12.dist-info/METADATA,sha256=TvTeDtkVFGsyMPWBgh3ANr5eKM1bLSi6SgI1hv4TUqw,9256
-dissect_cstruct-5.0.dev12.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
-dissect_cstruct-5.0.dev12.dist-info/entry_points.txt,sha256=z53zqZqwD2OLrAkRwrP4wTeiU9CQe7xrMly0T2c0_wQ,71
-dissect_cstruct-5.0.dev12.dist-info/top_level.txt,sha256=Mn-CQzEYsAbkxrUI0TnplHuXnGVKzxpDw_po_sXpvv4,8
-dissect_cstruct-5.0.dev12.dist-info/RECORD,,
+dissect/cstruct/types/wchar.py,sha256=VrnPXsP_lsyA-1xr4D2aSIVijumMhUuZWNlDHcn12-s,3964
+dissect_cstruct-5.0.dev14.dist-info/licenses/COPYRIGHT,sha256=H-18RXfshdH9AdHwW2eO1Xa-5s6tY5eipHh5c0whDu4,316
+dissect_cstruct-5.0.dev14.dist-info/licenses/LICENSE,sha256=PhUqiw6jAh2KbBdVRPBq_hfAvfcTBin7nZ3CK7NQbTM,11341
+dissect_cstruct-5.0.dev14.dist-info/METADATA,sha256=0EhDz7fNITZQWzSYoy46-j0nObu4B1Am3-nEipozmxc,9256
+dissect_cstruct-5.0.dev14.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
+dissect_cstruct-5.0.dev14.dist-info/entry_points.txt,sha256=z53zqZqwD2OLrAkRwrP4wTeiU9CQe7xrMly0T2c0_wQ,71
+dissect_cstruct-5.0.dev14.dist-info/top_level.txt,sha256=Mn-CQzEYsAbkxrUI0TnplHuXnGVKzxpDw_po_sXpvv4,8
+dissect_cstruct-5.0.dev14.dist-info/RECORD,,