fix(compiler): preserve proto package info in IR type lookup and validation

Signed-off-by: Peiyang He <peiyang_he@smail.nju.edu.cn>

Author: BaldDemian

compiler/fory_compiler/frontend/proto/translator.py

@@ -143,6 +143,7 @@ def _translate_enum(self, proto_enum: ProtoEnum) -> Enum:
             line=proto_enum.line,
             column=proto_enum.column,
             location=self._location(proto_enum.line, proto_enum.column),
+            package=self.proto_schema.package,
         )
 
     def _translate_message(self, proto_msg: ProtoMessage) -> Message:
@@ -173,6 +174,7 @@ def _translate_message(self, proto_msg: ProtoMessage) -> Message:
             line=proto_msg.line,
             column=proto_msg.column,
             location=self._location(proto_msg.line, proto_msg.column),
+            package=self.proto_schema.package,
         )
 
     def _translate_field(self, proto_field: ProtoField) -> Field:
@@ -261,6 +263,7 @@ def _translate_oneof(
             line=oneof.line,
             column=oneof.column,
             location=self._location(oneof.line, oneof.column),
+            package=self.proto_schema.package,
         )
 
     def _translate_oneof_case(self, proto_field: ProtoField) -> Field:

compiler/fory_compiler/ir/ast.py

@@ -178,6 +178,7 @@ class Message:
     location: Optional[SourceLocation] = None
     id_generated: bool = False
     id_source: Optional[str] = None
+    package: Optional[str] = None  # None if not protobuf
 
     def __repr__(self) -> str:
         id_str = f" [id={self.type_id}]" if self.type_id is not None else ""
@@ -218,6 +219,7 @@ class Enum:
     location: Optional[SourceLocation] = None
     id_generated: bool = False
     id_source: Optional[str] = None
+    package: Optional[str] = None  # None if not protobuf
 
     def __repr__(self) -> str:
         id_str = f" [id={self.type_id}]" if self.type_id is not None else ""
@@ -238,6 +240,7 @@ class Union:
     location: Optional[SourceLocation] = None
     id_generated: bool = False
     id_source: Optional[str] = None
+    package: Optional[str] = None  # None if not protobuf
 
     def __repr__(self) -> str:
         id_str = f" [id={self.type_id}]" if self.type_id is not None else ""
@@ -317,16 +320,26 @@ def get_option(self, name: str, default: Optional[str] = None) -> Optional[str]:
         return self.options.get(name, default)
 
     def get_type(self, name: str) -> Optional[TypingUnion[Message, Enum, "Union"]]:
-        """Look up a type by name, supporting qualified names like Parent.Child."""
-        # Handle qualified names (e.g., SearchResponse.Result)
-        if "." in name:
-            parts = name.split(".")
-            # Find the top-level type
+        """Look up a type by name, supporting qualified names like Parent.Child,
+        package.TypeName, and .package.TypeName (absolute proto-style)."""
+        cleaned = name.lstrip(".")
+        if "." in cleaned:
+            parts = cleaned.split(".")
+            # Try the first component as a top-level type name (e.g. "A" in "A.X").
             current = self._get_top_level_type(parts[0])
+            if current is None:
+                # First component not found as a type, then treat it as a package prefix.
+                # (e.g. "demo" in "demo.Foo" or "demo" in "demo.A.X").
+                if len(parts) >= 2:
+                    current = self._get_top_level_type_in_package(parts[1], parts[0])
+                    remaining = parts[2:]
+                else:
+                    return None
+            else:
+                remaining = parts[1:]
             if current is None:
                 return None
-            # Navigate through nested types
-            for part in parts[1:]:
+            for part in remaining:
                 if isinstance(current, Message):
                     current = current.get_nested_type(part)
                     if current is None:
@@ -336,7 +349,7 @@ def get_type(self, name: str) -> Optional[TypingUnion[Message, Enum, "Union"]]:
                     return None
             return current
         else:
-            return self._get_top_level_type(name)
+            return self._get_top_level_type(cleaned)
 
     def _get_top_level_type(
         self, name: str
@@ -353,6 +366,21 @@ def _get_top_level_type(
                 return message
         return None
 
+    def _get_top_level_type_in_package(
+        self, name: str, package: str
+    ) -> Optional[TypingUnion[Message, Enum, "Union"]]:
+        """Look up a top-level type by simple name within a specific package."""
+        for enum in self.enums:
+            if enum.name == name and enum.package == package:
+                return enum
+        for union in self.unions:
+            if union.name == name and union.package == package:
+                return union
+        for message in self.messages:
+            if message.name == name and message.package == package:
+                return message
+        return None
+
     def get_all_types(self) -> List[TypingUnion[Message, Enum, "Union"]]:
         """Get all types including nested types (flattened)."""
         result: List[TypingUnion[Message, Enum, "Union"]] = []

compiler/fory_compiler/ir/validator.py

@@ -86,18 +86,20 @@ def _apply_type_id_defaults(self) -> None:
             if t.type_id is not None:
                 used_ids[t.type_id] = t
 
-        def qualify(full_name: str) -> str:
-            package = self.schema.package_alias or self.schema.package
+        def qualify(full_name: str, type_package: Optional[str] = None) -> str:
+            package = type_package or self.schema.package_alias or self.schema.package
             if package:
                 return f"{package}.{full_name}"
             return full_name
 
-        def resolve_hash_source(full_name: str, alias: Optional[str]) -> str:
+        def resolve_hash_source(
+            full_name: str, alias: Optional[str], type_package: Optional[str] = None
+        ) -> str:
             if not alias:
-                return qualify(full_name)
+                return qualify(full_name, type_package)
             if "." in alias:
                 return alias
-            package = self.schema.package_alias or self.schema.package
+            package = type_package or self.schema.package_alias or self.schema.package
             if package:
                 return f"{package}.{alias}"
             return alias
@@ -106,7 +108,8 @@ def assign_id(type_def, full_name: str) -> None:
             if type_def.type_id is not None:
                 return
             alias = type_def.options.get("alias")
-            source_name = resolve_hash_source(full_name, alias)
+            type_package = getattr(type_def, "package", None)
+            source_name = resolve_hash_source(full_name, alias, type_package)
             generated_id = compute_registered_type_id(source_name)
             if generated_id in used_ids:
                 self._error(
@@ -144,28 +147,33 @@ def walk_message(message: Message, parent_path: str = "") -> None:
             walk_message(message)
 
     def _check_duplicate_type_names(self) -> None:
+        # Define (name, package) as key to detect duplication so that same-named types in different packages are not flagged.
+        # (package is None if not protobuf)
         names = {}
         for enum in self.schema.enums:
-            if enum.name in names:
+            key = (enum.name, getattr(enum, "package", None))
+            if key in names:
                 self._error(
                     f"Duplicate type name: {enum.name}",
-                    enum.location or names[enum.name],
+                    enum.location or names[key],
                 )
-            names.setdefault(enum.name, enum.location)
+            names.setdefault(key, enum.location)
         for union in self.schema.unions:
-            if union.name in names:
+            key = (union.name, getattr(union, "package", None))
+            if key in names:
                 self._error(
                     f"Duplicate type name: {union.name}",
-                    union.location or names[union.name],
+                    union.location or names[key],
                 )
-            names.setdefault(union.name, union.location)
+            names.setdefault(key, union.location)
         for message in self.schema.messages:
-            if message.name in names:
+            key = (message.name, getattr(message, "package", None))
+            if key in names:
                 self._error(
                     f"Duplicate type name: {message.name}",
-                    message.location or names[message.name],
+                    message.location or names[key],
                 )
-            names.setdefault(message.name, message.location)
+            names.setdefault(key, message.location)
 
     def _check_duplicate_type_ids(self) -> None:
         type_ids = {}
@@ -309,15 +317,12 @@ def _is_message_type(
     def _resolve_named_type(
         self, name: str, parent_stack: List[Message]
     ) -> Optional[TypingUnion[Message, Enum, Union]]:
-        parts = name.split(".")
+        cleaned = name.lstrip(".")
+        parts = cleaned.split(".")
         if len(parts) > 1:
-            current = self._find_top_level_type(parts[0])
-            for part in parts[1:]:
-                if isinstance(current, Message):
-                    current = current.get_nested_type(part)
-                else:
-                    return None
-            return current
+            # Call schema.get_type to handle both nested-type paths
+            # (e.g. A.X) and package-prefixed paths (e.g. demo.Foo, demo.A.X).
+            return self.schema.get_type(name)
         for msg in reversed(parent_stack):
             nested = msg.get_nested_type(name)
             if nested is not None:
@@ -354,8 +359,28 @@ def check_type_ref(
                             found = True
                             break
 
-                if not found and self.schema.get_type(type_name) is not None:
-                    found = True
+                if not found:
+                    resolved = self.schema.get_type(type_name)
+                    if resolved is not None:
+                        found = True
+                        # Handle case like test_proto_imported_package_qualified_types_fail properly.
+                        # Reject unqualified type names from a different package.
+                        # This check only happens when both type and schema have explicit package info,
+                        # so that non-protobuf schemas whose package is None are unaffected.
+                        if "." not in type_name:
+                            type_package = getattr(resolved, "package", None)
+                            schema_package = self.schema.package
+                            if (
+                                type_package is not None
+                                and schema_package is not None
+                                and type_package != schema_package
+                            ):
+                                self._error(
+                                    f"Type '{type_name}' belongs to package "
+                                    f"'{type_package}'; use '{type_package}.{type_name}' "
+                                    f"to reference it from package '{schema_package}'",
+                                    field.location,
+                                )
 
                 if not found:
                     self._error(f"Unknown type '{type_name}'", field.location)

compiler/fory_compiler/tests/test_proto_frontend.py

@@ -20,6 +20,10 @@
 from fory_compiler.frontend.proto import ProtoFrontend
 from fory_compiler.ir.ast import PrimitiveType
 from fory_compiler.ir.types import PrimitiveKind
+from fory_compiler.cli import resolve_imports
+from fory_compiler.ir.validator import SchemaValidator
+import tempfile
+from pathlib import Path
 
 
 def test_proto_type_mapping():
@@ -106,3 +110,126 @@ def test_proto_file_option_enable_auto_type_id():
     """
     schema = ProtoFrontend().parse(source)
     assert schema.get_option("enable_auto_type_id") is False
+
+
+def test_proto_nested_qualified_types_pass():
+    source = """
+    syntax = "proto3";
+    package demo;
+
+    message A {
+        message X{}
+    }
+    message B {
+        A.X x1 = 1;
+        .demo.A.X x2 = 2;
+        demo.A.X x3 = 3;
+    }
+    """
+    schema = ProtoFrontend().parse(source)
+    validator = SchemaValidator(schema)
+    assert validator.validate()
+
+
+def test_proto_nested_qualified_types_fail():
+    source = """
+    syntax = "proto3";
+    package demo;
+
+    message A {
+        message X{}
+    }
+    message B {
+        X x = 1;
+    }
+    """
+    schema = ProtoFrontend().parse(source)
+    validator = SchemaValidator(schema)
+    assert not validator.validate()
+
+
+def test_proto_same_package_qualified_types_pass():
+    source = """
+    syntax = "proto3";
+    package demo;
+
+    message Foo {}
+
+    message Bar {
+        demo.Foo  foo1 = 1;
+        .demo.Foo foo2 = 2;
+    }
+    """
+    schema = ProtoFrontend().parse(source)
+    validator = SchemaValidator(schema)
+    assert validator.validate()
+
+
+def test_proto_imported_package_qualified_types_fail():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tmpdir = Path(tmpdir)
+        common_proto = tmpdir / "common.proto"
+        common_proto.write_text(
+            """
+            syntax = "proto3";
+            package common;
+
+            message Address {}
+            """
+        )
+        main_proto = tmpdir / "main.proto"
+        main_proto.write_text(
+            """
+            syntax = "proto3";
+            package main;
+            import "common.proto";
+
+            message User {
+                Address addr1 = 1;
+                Address addr2 = 2;
+            }
+            """
+        )
+        schema = resolve_imports(main_proto, [tmpdir])
+        validator = SchemaValidator(schema)
+        assert not validator.validate()
+
+
+def test_proto_imported_package_qualified_types_pass():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tmpdir = Path(tmpdir)
+        common1_proto = tmpdir / "common1.proto"
+        common1_proto.write_text(
+            """
+            syntax = "proto3";
+            package common1;
+
+            message Address {}
+            """
+        )
+        common2_proto = tmpdir / "common2.proto"
+        common2_proto.write_text(
+            """
+            syntax = "proto3";
+            package common2;
+
+            message Address {}
+            """
+        )
+        main_proto = tmpdir / "main.proto"
+        main_proto.write_text(
+            """
+            syntax = "proto3";
+            package main;
+            import "common1.proto";
+            import "common2.proto";
+
+            message User {
+                common1.Address addr1 = 1;
+                .common2.Address addr2 = 2;
+            }
+            """
+        )
+        schema = resolve_imports(main_proto, [tmpdir])
+        validator = SchemaValidator(schema)
+        assert validator.validate()