Process feedback

Co-authored-by: Copilot <copilot@github.com>

Author: Schamper

dissect/cstruct/expression.py

@@ -3,7 +3,7 @@
 from typing import TYPE_CHECKING
 
 from dissect.cstruct.exceptions import ExpressionParserError
-from dissect.cstruct.lexer import _IDENTIFIER_TYPES, Lexer, TokenCursor, TokenType
+from dissect.cstruct.lexer import IDENTIFIER_TYPES, Lexer, TokenCursor, TokenType
 from dissect.cstruct.utils import offsetof, sizeof
 
 if TYPE_CHECKING:
@@ -12,7 +12,6 @@
     from dissect.cstruct import cstruct
     from dissect.cstruct.lexer import Token
 
-
 BINARY_OPERATORS: dict[TokenType, Callable[[int, int], int]] = {
     TokenType.PIPE: lambda a, b: a | b,
     TokenType.CARET: lambda a, b: a ^ b,
@@ -155,15 +154,32 @@ def evaluate(self, cs: cstruct, context: dict[str, int] | None = None) -> int:
                 # Evaluate immediately
                 self._evaluate_expression(cs)
 
-            elif token.type in _IDENTIFIER_TYPES:
-                if token.value in context:
-                    self._queue.append(int(context[self._take().value]))
-
-                elif token.value in cs.consts:
-                    self._queue.append(int(cs.consts[self._take().value]))
-
-                else:
-                    raise self._error(f"Unknown identifier: '{token.value}'", token=token)
+            elif token.type in IDENTIFIER_TYPES:
+                ident = self._take().value
+
+                obj = None
+                for i, part in enumerate(ident.split(".")):
+                    if i == 0:
+                        if part in context:
+                            obj = context[part]
+                        elif part in cs.consts:
+                            obj = cs.consts[part]
+                        elif part in cs.typedefs:
+                            obj = cs.resolve(part)
+                        else:
+                            raise self._error(f"Unknown identifier: '{ident}'", token=token)
+                    else:
+                        if isinstance(obj, dict) and part in obj:
+                            obj = obj[part]
+                        elif hasattr(obj, part):
+                            obj = getattr(obj, part)
+                        else:
+                            raise self._error(f"Unknown identifier: '{ident}'", token=token)
+
+                try:
+                    self._queue.append(int(obj))
+                except (ValueError, TypeError):
+                    raise self._error(f"Identifier '{ident}' does not resolve to an integer", token=token)
 
             elif token.type == TokenType.LPAREN:
                 if self._previous().type == TokenType.NUMBER:

dissect/cstruct/lexer.py

@@ -100,8 +100,6 @@ def __repr__(self) -> str:
     "typedef": TokenType.TYPEDEF,
 }
 
-_IDENTIFIER_TYPES = set(_C_KEYWORDS.values()) | {TokenType.IDENTIFIER}
-
 _SINGLE_CHARS = {
     "{": TokenType.LBRACE,
     "}": TokenType.RBRACE,
@@ -125,9 +123,11 @@ def __repr__(self) -> str:
     "~": TokenType.TILDE,
 }
 
-_RE_IDENTIFIER = re.compile(r"[a-zA-Z_][a-zA-Z0-9_]*")
+_RE_IDENTIFIER = re.compile(r"[a-zA-Z_.][a-zA-Z0-9_.]*")
 _RE_WHITESPACE = re.compile(r"[ \t\r\n]+")
 
+IDENTIFIER_TYPES = set(_C_KEYWORDS.values()) | {TokenType.IDENTIFIER}
+
 
 def tokenize(data: str) -> list[Token]:
     """Convenience function to tokenize input data."""

dissect/cstruct/parser.py

@@ -11,7 +11,7 @@
     ParserError,
 )
 from dissect.cstruct.expression import Expression
-from dissect.cstruct.lexer import _IDENTIFIER_TYPES, TokenCursor, TokenType, tokenize
+from dissect.cstruct.lexer import IDENTIFIER_TYPES, TokenCursor, TokenType, tokenize
 from dissect.cstruct.types import BaseArray, BaseType, Field, Structure
 
 if TYPE_CHECKING:
@@ -74,12 +74,7 @@ def parse(self, data: str) -> None:
 
         data = _join_line_continuations(data)
 
-        # Tokenize and preprocess the input, then parse top-level definitions
         self._reset_tokens(tokenize(data))
-        preprocessed_tokens = self._preprocess()
-        self.reset()
-
-        self._reset_tokens(preprocessed_tokens)
         self._parse()
 
     def _match(self, *types: TokenType) -> Token | None:
@@ -95,46 +90,41 @@ def _at(self, *types: TokenType) -> bool:
     def _error(self, msg: str, *, token: Token | None = None) -> ParserError:
         return ParserError(f"line {(token if token is not None else self._tokens[self._pos]).line}: {msg}")
 
-    def _preprocess(self) -> list[Token]:
-        """Handle preprocessor directives and return a new list of tokens with directives processed."""
-        result = []
+    def _in_false_branch(self) -> bool:
+        """Return whether we're currently in a false conditional branch."""
+        return bool(self._conditional_stack) and not self._conditional_stack[-1][1]
 
-        while self._tokens[self._pos].type != TokenType.EOF:
-            token = self._tokens[self._pos]
+    def _handle_preprocessor(self) -> bool:
+        """Handle preprocessor directives.
 
-            # Always handle conditional directives first (even in false branches)
-            if token.type in (TokenType.PP_IFDEF, TokenType.PP_IFNDEF, TokenType.PP_ELSE, TokenType.PP_ENDIF):
-                self._handle_conditional()
-                continue
+        Returns ``True`` if a directive was processed, indicating the caller should continue its loop.
+        """
+        token = self._current()
 
-            # If we're in a false conditional branch, skip this token
-            if self._conditional_stack and not self._conditional_stack[-1][1]:
-                self._pos += 1
-                continue
+        if token.type in (TokenType.PP_IFDEF, TokenType.PP_IFNDEF, TokenType.PP_ELSE, TokenType.PP_ENDIF):
+            self._handle_conditional()
+            return True
 
-            if token.type == TokenType.PP_DEFINE:
-                self._parse_define()
-            elif token.type == TokenType.PP_UNDEF:
-                self._parse_undef()
-            elif token.type == TokenType.PP_INCLUDE:
-                self._parse_include()
-            else:
-                # Not a preprocessor directive, just add it to the result
-                result.append(token)
-                self._pos += 1
+        if token.type == TokenType.PP_DEFINE:
+            self._parse_define()
+            return True
 
-        # Append EOF token
-        result.append(self._tokens[self._pos])
-        self._pos += 1
+        if token.type == TokenType.PP_UNDEF:
+            self._parse_undef()
+            return True
 
-        if self._conditional_stack:
-            raise self._error("unclosed conditional statement", token=self._conditional_stack[-1][0])
+        if token.type == TokenType.PP_INCLUDE:
+            self._parse_include()
+            return True
 
-        return result
+        return False
 
     def _parse(self) -> None:
         """Parse top-level definitions from the token stream."""
         while (token := self._current()).type != TokenType.EOF:
+            if self._handle_preprocessor():
+                continue
+
             if token.type == TokenType.PP_FLAGS:
                 self._parse_config_flags()
             elif token.type == TokenType.TYPEDEF:
@@ -158,6 +148,9 @@ def _parse(self) -> None:
             else:
                 raise self._error(f"unexpected token {token.value!r}")
 
+        if self._conditional_stack:
+            raise self._error("unclosed conditional statement", token=self._conditional_stack[-1][0])
+
     # Preprocessor directives
 
     def _parse_define(self) -> None:
@@ -246,6 +239,34 @@ def _handle_conditional(self) -> None:
                 raise self._error("#endif without matching #ifdef/#ifndef", token=token)
             self._conditional_stack.pop()
 
+        # If we ended up in a false branch, skip ahead to the matching #else/#endif
+        if self._in_false_branch():
+            self._skip_false_branch()
+
+    def _skip_false_branch(self) -> None:
+        """Skip all tokens in a false conditional branch until the matching ``#else`` or ``#endif``."""
+        depth = 0
+        while self._current().type != TokenType.EOF:
+            token_type = self._current().type
+
+            if token_type in (TokenType.PP_IFDEF, TokenType.PP_IFNDEF):
+                depth += 1
+                self._pos += 1
+                # Skip the identifier argument
+                if self._current().type != TokenType.EOF:
+                    self._pos += 1
+            elif token_type == TokenType.PP_ENDIF:
+                if depth == 0:
+                    return
+                depth -= 1
+                self._pos += 1
+            elif token_type == TokenType.PP_ELSE:
+                if depth == 0:
+                    return
+                self._pos += 1
+            else:
+                self._pos += 1
+
     # Type definitions
 
     def _parse_typedef(self) -> None:
@@ -347,6 +368,8 @@ def _parse_enum_or_flag(self) -> type[Enum | Flag]:
         values: dict[str, int] = {}
 
         while not self._at(TokenType.RBRACE):
+            if self._handle_preprocessor():
+                continue
             self._assert_not_eof()
 
             member_name = self._expect(TokenType.IDENTIFIER).value
@@ -385,6 +408,8 @@ def _parse_field_list(self) -> list[Field]:
         fields: list[Field] = []
 
         while not self._at(TokenType.RBRACE):
+            if self._handle_preprocessor():
+                continue
             self._assert_not_eof()
 
             fields.append(self._parse_field())
@@ -419,7 +444,7 @@ def _parse_field_name(self, base_type: type[BaseType]) -> tuple[type[BaseType],
             type_ = self.cs._make_pointer(type_)
 
         # Field name
-        name = self._expect(*_IDENTIFIER_TYPES).value
+        name = self._expect(*IDENTIFIER_TYPES).value
 
         # Array dimensions
         type_ = self._parse_array_dimensions(type_)

tests/test_benchmark.py

@@ -187,9 +187,6 @@ def test_benchmark_parser(cs: cstruct, benchmark: BenchmarkFixture) -> None:
     tokens = tokenize(cdef)
 
     def parse(parser: CStyleParser, tokens: list[Token]) -> None:
-        parser.reset()
-        parser._tokens = tokens
-        tokens = parser._preprocess()
         parser.reset()
         parser._tokens = tokens
         parser._parse()

tests/test_expression.py

@@ -167,9 +167,33 @@ def test_expression(cs_with_consts: cstruct, expression: str, answer: int) -> No
         pytest.param(
             "sizeof(0 +)",
             ResolveError,
-            "Unknown type 0 +",
+            "Unknown type 0 \\+",
             id="invalid-sizeof-expression",
         ),
+        pytest.param(
+            "UNKNOWN",
+            ExpressionParserError,
+            "Unknown identifier: 'UNKNOWN'",
+            id="unknown-identifier",
+        ),
+        pytest.param(
+            "UNKNOWN.UNKNOWN",
+            ExpressionParserError,
+            "Unknown identifier: 'UNKNOWN.UNKNOWN'",
+            id="unknown-identifier",
+        ),
+        pytest.param(
+            "1 2",
+            ExpressionParserError,
+            "Invalid expression: too many operands",
+            id="too-many-operands",
+        ),
+        pytest.param(
+            "1 + + 2",
+            ExpressionParserError,
+            "Invalid expression: not enough operands",
+            id="double-binary-operator",
+        ),
     ],
 )
 def test_expression_failure(cs_with_consts: cstruct, expression: str, exception: type, message: str) -> None:
@@ -210,3 +234,38 @@ def test_offsetof(cs: cstruct) -> None:
     assert Expression("offsetof(test, b)").evaluate(cs) == 4
     assert Expression("offsetof(test, c)").evaluate(cs) == 12
     assert Expression("offsetof(test, d)").evaluate(cs) == 14
+
+
+def test_dotted_identifier(cs: cstruct) -> None:
+    """Test dotted member access in expressions."""
+    cs.load("""
+    enum color : uint8 {
+        RED = 1,
+        GREEN = 2,
+        BLUE = 4
+    };
+    """)
+
+    assert Expression("color.RED").evaluate(cs) == 1
+    assert Expression("color.RED | color.BLUE").evaluate(cs) == 5
+    assert Expression("color.GREEN + 10").evaluate(cs) == 12
+
+
+def test_dotted_identifier_unknown_member(cs: cstruct) -> None:
+    """Test error when a member in a dotted expression doesn't exist."""
+    cs.load("""
+    enum color : uint8 {
+        RED = 1
+    };
+    """)
+
+    with pytest.raises(ExpressionParserError, match="Unknown identifier: 'color\\.NOTAMEMBER'"):
+        Expression("color.NOTAMEMBER").evaluate(cs)
+
+
+def test_dotted_identifier_not_integer(cs: cstruct) -> None:
+    """Test error when a dotted expression resolves to a non-integer."""
+    cs.consts["mydict"] = {"nested": "hello"}
+
+    with pytest.raises(ExpressionParserError, match="does not resolve to an integer"):
+        Expression("mydict.nested").evaluate(cs)