Add method to dump C definition (#160)

Author: Schamper

dissect/cstruct/cstruct.py

@@ -302,6 +302,42 @@ def loadfile(self, path: str, deftype: int | None = None, **kwargs) -> None:
         with Path(path).open() as fh:
             self.load(fh.read(), deftype, **kwargs)
 
+    def cdef(self) -> str:
+        """Render all constants, structure, union, enum and flag definitions back to their C-style definitions.
+
+        Note that the result is semantically equivalent to the original definitions, but not necessarily identical.
+
+        Returns:
+            The C-style definitions as a string.
+        """
+        empty = cstruct()
+
+        blocks = []
+
+        defines = []
+        for name, value in self.consts.items():
+            if name in empty.consts:
+                continue
+
+            defines.append(f"#define {name} {value!r}")
+
+        if defines:
+            blocks.append("\n".join(defines))
+
+        for name, typedef in self.typedefs.items():
+            if name in empty.typedefs or not isinstance(typedef, type):
+                continue
+
+            if not issubclass(typedef, (Structure, Enum, Flag)):
+                continue
+
+            if typedef.__name__ == name:
+                blocks.append(typedef.cdef())
+            else:
+                blocks.append(f"typedef {typedef.__name__} {name};")
+
+        return "\n\n".join(blocks)
+
     def read(self, name: str, stream: BinaryIO) -> Any:
         """Parse data using a given type.
 

dissect/cstruct/types/enum.py

@@ -104,6 +104,24 @@ def _write_0(cls, stream: BinaryIO, array: list[BaseType | int]) -> int:
         data = [entry.value if isinstance(entry, _Enum) else entry for entry in array]
         return cls._write_array(stream, [*data, cls.type.__default__()])
 
+    def cdef(cls) -> str:
+        """Render this enum (or flag) back to its C-style definition.
+
+        Returns:
+            The C-style definition as a string.
+        """
+        keyword = "flag" if issubclass(cls, IntFlag) else "enum"
+        underlying = cls.type.__name__
+        title = f"{keyword} {cls.__name__} : {underlying}" if cls.__name__ else f"{keyword} : {underlying}"
+
+        lines = [f"{title} {{"]
+        for member_name, member in cls.__members__.items():
+            value = hex(member.value) if keyword == "flag" else member.value
+            lines.append(f"    {member_name} = {value},")
+        lines.append("};")
+
+        return "\n".join(lines)
+
 
 def _fix_alias_members(cls: type[Enum]) -> None:
     # Emulate aenum NoAlias behaviour

dissect/cstruct/types/structure.py

@@ -13,7 +13,10 @@
 from typing import TYPE_CHECKING, Any, BinaryIO
 
 from dissect.cstruct.bitbuffer import BitBuffer
+from dissect.cstruct.expression import Expression
+from dissect.cstruct.lexer import IDENTIFIER_TYPES
 from dissect.cstruct.types.base import (
+    BaseArray,
     BaseType,
     MetaType,
     _is_buffer_type,
@@ -380,6 +383,102 @@ def commit(cls) -> None:
         for key, value in classdict.items():
             setattr(cls, key, value)
 
+    def cdef(cls, *, recursive: bool = False) -> str:
+        """Render this structure (or union) back to its C-style definition.
+
+        When ``recursive=True``, any ``#define`` constants referenced in array sizes are emitted as well, so that the
+        result can be parsed on its own. Note that the result is semantically equivalent to the original definition,
+        but not necessarily identical.
+
+        Args:
+            recursive: Also emit the definitions of any types referenced by this structure.
+
+        Returns:
+            The C-style definition as a string.
+        """
+
+        def _decompose(type_: MetaType) -> tuple[MetaType, int, list[int | Expression | None]]:
+            dimensions = []
+            while issubclass(type_, BaseArray):
+                dimensions.append(type_.num_entries)
+                type_ = type_.type
+
+            stars = 0
+            while issubclass(type_, Pointer):
+                stars += 1
+                type_ = type_.type
+
+            return type_, stars, dimensions
+
+        def _deps(type_: StructureMetaType, seen: set[MetaType]) -> list[MetaType | str]:
+            deps = []
+            for field in type_.__fields__:
+                base, _, dimensions = _decompose(field.type)
+
+                for dim in dimensions:
+                    if isinstance(dim, Expression):
+                        for token in dim._tokens:
+                            if (
+                                token.type in IDENTIFIER_TYPES
+                                and token.value in cls.cs.consts
+                                and token.value not in seen
+                            ):
+                                deps.append(token.value)
+                                seen.add(token.value)
+
+                if isinstance(base, StructureMetaType) and base not in seen:
+                    if base.__anonymous__:
+                        # Anonymous structures are not emitted, but their dependencies are
+                        deps.extend(_deps(base, seen))
+                    else:
+                        seen.add(base)
+                        deps.extend(_deps(base, seen))
+                        deps.append(base)
+
+                elif isinstance(base, EnumMetaType) and base not in seen:
+                    seen.add(base)
+                    deps.append(base)
+
+            return deps
+
+        def _render_struct(type_: StructureMetaType, name: str) -> list[str]:
+            keyword = "union" if isinstance(type_, UnionMetaType) else "struct"
+            title = f"{keyword} {name}" if name else keyword
+
+            lines = [f"{title} {{"]
+            for field in type_.__fields__:
+                base, stars, dimensions = _decompose(field.type)
+
+                declarator = "*" * stars + (field.name or "") + "".join(f"[{dim or ''}]" for dim in dimensions)
+                if field.bits is not None:
+                    declarator = f"{declarator} : {field.bits}" if declarator else f": {field.bits}"
+
+                if isinstance(base, StructureMetaType) and base.__anonymous__:
+                    # Anonymous struct/union: inline the full definition
+                    lines.extend(f"    {line}" for line in _render_struct(base, ""))
+                    lines[-1] = f"{lines[-1]} {declarator};" if declarator else f"{lines[-1]};"
+                else:
+                    prefix = f"{base.__name__} " if declarator else base.__name__
+                    lines.append(f"    {prefix}{declarator};")
+
+            lines.append("}")
+            return lines
+
+        deps = _deps(cls, {cls}) if recursive else []
+        deps.append(cls)
+
+        consts = []
+        blocks = []
+        for type_ in deps:
+            if isinstance(type_, str):
+                consts.append(f"#define {type_} {cls.cs.consts[type_]!r}")
+            elif isinstance(type_, EnumMetaType):
+                blocks.append(type_.cdef())
+            else:
+                blocks.append("\n".join(_render_struct(type_, type_.__name__)) + ";")
+
+        return ("\n".join(consts) + "\n\n" if consts else "") + ("\n\n".join(blocks) if blocks else "")
+
 
 class Structure(BaseType, metaclass=StructureMetaType):
     """Base class for cstruct structure type classes.

tests/test_basic.py

@@ -551,3 +551,86 @@ def test_copy(cs: cstruct) -> None:
     # Verify that types in the copied cstruct reference the copied cstruct
     assert cs_copy.resolve("test").cs is cs_copy
     assert cs_copy.resolve("test2").cs is cs_copy
+
+
+def test_cdef(cs: cstruct) -> None:
+    """Test that a cstruct instance can render all its definitions back to C-style definitions."""
+    cs.load(
+        """
+        #define MAGIC_LENGTH 4
+        enum Color : uint16 { RED = 1, GREEN = 2, BLUE = 4 };
+        struct Point { int32 x; int32 y; };
+        union Val { uint32 as_int; float as_float; };
+        typedef Point Coord;
+        struct Header {
+            char magic[MAGIC_LENGTH];
+            Color color;
+            Point points[2];
+            Val v;
+        };
+        """
+    )
+
+    assert cs.cdef() == textwrap.dedent(
+        """\
+        #define MAGIC_LENGTH 4
+
+        enum Color : uint16 {
+            RED = 1,
+            GREEN = 2,
+            BLUE = 4,
+        };
+
+        struct Point {
+            int32 x;
+            int32 y;
+        };
+
+        union Val {
+            uint32 as_int;
+            float as_float;
+        };
+
+        typedef Point Coord;
+
+        struct Header {
+            char magic[4];
+            Color color;
+            Point points[2];
+            Val v;
+        };"""
+    )
+
+    # The full dump should be parseable on its own into equivalent types
+    reparsed = cstruct()
+    reparsed.load(cs.cdef())
+    assert reparsed.consts == cs.consts
+    for name in ("Color", "Point", "Val", "Coord", "Header"):
+        assert getattr(cs, name).size == getattr(reparsed, name).size
+
+
+def test_cdef_defines(cs: cstruct) -> None:
+    """Test that ``#define`` constants (including those referenced in expressions) are dumped and round-trip."""
+    cs.load(
+        """
+        #define INT_CONST 8
+        #define STR_CONST "hello"
+        #define BYTES_CONST b"world"
+
+        struct Test {
+            uint32 count;
+            uint8 data[count * INT_CONST];
+        };
+        """
+    )
+
+    src = cs.cdef()
+    assert "#define INT_CONST 8" in src
+    assert "#define STR_CONST 'hello'" in src
+    assert "#define BYTES_CONST b'world'" in src
+    # The constant referenced in the (dynamic) array expression must be emitted
+    assert "uint8 data[count * INT_CONST];" in src
+
+    reparsed = cstruct()
+    reparsed.load(src)
+    assert reparsed.consts == cs.consts

tests/test_types_enum.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from enum import Enum as StdEnum
+from textwrap import dedent
 from typing import TYPE_CHECKING
 
 import pytest
@@ -448,3 +449,16 @@ def test_enum_default_default(cs: cstruct) -> None:
     cs.load(cdef)
 
     assert cs.test.__default__() == cs.test.default == cs.test(0)
+
+
+def test_cdef_enum(cs: cstruct) -> None:
+    cs.load("enum Color : uint16 { RED = 1, GREEN = 2, BLUE = 4 };")
+
+    assert cs.Color.cdef() == dedent(
+        """\
+        enum Color : uint16 {
+            RED = 1,
+            GREEN = 2,
+            BLUE = 4,
+        };"""
+    )

tests/test_types_flag.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from enum import Flag as StdFlag
+from textwrap import dedent
 from typing import TYPE_CHECKING
 
 import pytest
@@ -282,3 +283,16 @@ def test_flag_default(cs: cstruct) -> None:
     assert cs.test.__default__() == cs.test(0)
     assert cs.test[1].__default__() == [cs.test(0)]
     assert cs.test[None].__default__() == []
+
+
+def test_cdef_flag(cs: cstruct) -> None:
+    cs.load("flag Perm : uint8 { R, W, X };")
+
+    assert cs.Perm.cdef() == dedent(
+        """\
+        flag Perm : uint8 {
+            R = 0x1,
+            W = 0x2,
+            X = 0x4,
+        };"""
+    )