Read a define value untill the end of the line or when it encounters a comment

dissect/cstruct/lexer.py

@@ -390,6 +390,52 @@ def _read_string(self) -> str:
 
         return result
 
+    def _read_define_value(self) -> str:
+        """Read value of a define, ignoring comments and unecessary white space."""
+        text_offsets: list[tuple[int, int]] = []
+
+        while not self.eof:
+            char = self._current()
+            if char == "\n":
+                break
+
+            # Read byte string
+            next = self._peek()
+            if char == "b" and next is not None and next in ("\"'"):
+                sof = self._pos
+                self._pos += 1  # consume the b
+                self._read_string()
+                text_offsets.append((sof, self._pos))
+                continue
+
+            # read normal string
+            if char in ("\"'"):
+                sof = self._pos
+                self._read_string()
+                text_offsets.append((sof, self._pos))
+                continue
+
+            # read word
+            if char not in " \t\r":
+                sof = self._pos
+                self._read_until(" \t\r\n")
+                text_offsets.append((sof, self._pos))
+                continue
+
+            # Skip whitespace
+            self._read_while(" \t\r", or_eof=True)
+            if self.eof:
+                break
+            self._skip_comment()
+            if self.eof:
+                break
+
+        result: list[str] = []
+        for start, end in text_offsets:
+            self._pos = start
+            result.append(self._take(end - start))
+        return " ".join(result)
+
     def _read_preprocessor(self) -> None:
         """Read a preprocessor directive starting with ``#``."""
         line = self._line
@@ -428,7 +474,7 @@ def _read_preprocessor(self) -> None:
                 # No value, just a simple macro definition
                 return
 
-            if (value := self._read_until("\n")).strip():
+            if value := self._read_define_value():
                 self._emit(TokenType.STRING, value, line)
 
         elif token_type == TokenType.PP_INCLUDE:

tests/test_lexer.py

@@ -96,7 +96,7 @@
             + 2)
             """,
             [TokenType.PP_DEFINE, TokenType.IDENTIFIER, TokenType.STRING],
-            ["define", "FOO", "(1             + 2)"],
+            ["define", "FOO", "(1 + 2)"],
         ),
         ("#undef", [TokenType.PP_UNDEF], ["undef"]),
         ("#ifdef", [TokenType.PP_IFDEF], ["ifdef"]),

tests/test_parser.py

@@ -251,11 +251,14 @@ def test_define(cs: cstruct) -> None:
     #define MULTILINE (1 + \
                         2 + \
                         3)
-    #define QUOTES "\'\"a'b\""
+    #define QUOTES "'\\"a'b\\""
     #define ESCAPE "\\'\\"a'b\\"\\n"
     #define BYTES_ESCAPE b"`\\n"
     #define FUNC(x) ( x == 0 )
     #define TERNARY(x) ( x ? 1 : 0 )
+    #define MULTI_LINE_QUOTE "\
+    #define DEFINITION "
+    #define SPACED_BYTES b"        "
     """
     cs.load(cdef)
 
@@ -275,6 +278,8 @@ def test_define(cs: cstruct) -> None:
     # We don't evaluate function-like macros yet, so they should be stored as their raw string representation
     assert cs.consts["FUNC"] == "(x) ( x == 0 )"
     assert cs.consts["TERNARY"] == "(x) ( x ? 1 : 0 )"
+    assert cs.consts["MULTI_LINE_QUOTE"] == "    #define DEFINITION "
+    assert cs.consts["SPACED_BYTES"] == b"        "
 
 
 def test_define_flag_value(cs: cstruct) -> None:
@@ -515,6 +520,34 @@ def test_preprocessor_in_struct_body(cs: cstruct) -> None:
     assert cs.test.fields["bonus"].type == cs.uint64
 
 
+def test_preprocessor_define_with_comments(cs: cstruct) -> None:
+    cdef = """
+    #define IGNORE_SCOPED_COMMENT       data1      /* ... */
+    #define IGNORE_LINE_COMMENT         data2      //
+    #define IGNORE_MULTILINE_COMMENT    data3      /* Multiline
+                                                      comment
+                                                    */
+    #define IGNORE_DEFINE_IN_COMMENT    data4      /* define
+    #define TEST_NO                     data_no       inside
+                                                      comment
+                                                    */
+    #define TEXT_CONTAINING_COMMENT     "text with comments /* data5 */"
+    #define IGNORE_STRING_AFTER_COMMENT "data6" // \
+                                        "string that cannot be reached"
+    #define NO_COMMENT_BETWEEN_CONCATINATION "data7" /* Comment */ \
+                                             "should get concatinated"
+    """
+
+    cs.load(cdef)
+    assert cs.consts["IGNORE_SCOPED_COMMENT"] == "data1"
+    assert cs.consts["IGNORE_LINE_COMMENT"] == "data2"
+    assert cs.consts["IGNORE_MULTILINE_COMMENT"] == "data3"
+    assert cs.consts["IGNORE_DEFINE_IN_COMMENT"] == "data4"
+    assert cs.consts["TEXT_CONTAINING_COMMENT"] == "text with comments /* data5 */"
+    assert cs.consts["IGNORE_STRING_AFTER_COMMENT"] == "data6"
+    assert cs.consts["NO_COMMENT_BETWEEN_CONCATINATION"] == 'data7" "should get concatinated'
+
+
 def test_preprocessor_define_from_enum_in_struct(cs: cstruct) -> None:
     """Test #define referencing enum values used for conditional fields and array sizes."""
     cdef = """