Merge pull request google#1888 from xlsynth:cdleary/2025-01-25-bad-ro…

…ot-type-info-switch

PiperOrigin-RevId: 720301920

xls/dslx/bytecode/bytecode_emitter.cc

@@ -797,11 +797,10 @@ absl::StatusOr<InterpValue> BytecodeEmitter::HandleColonRefInternal(
             return GetBuiltinNameDefColonAttr(builtin_name_def, node->attr());
           },
           [&](ArrayTypeAnnotation* array_type) -> absl::StatusOr<InterpValue> {
-            XLS_ASSIGN_OR_RETURN(
-                TypeInfo * type_info,
-                import_data_->GetRootTypeInfoForNode(array_type));
+            const TypeInfo& type_info = GetTypeInfoForNodeIfDifferentModule(
+                array_type, *type_info_, *import_data_);
             XLS_ASSIGN_OR_RETURN(InterpValue value,
-                                 type_info->GetConstExpr(array_type->dim()));
+                                 type_info.GetConstExpr(array_type->dim()));
             XLS_ASSIGN_OR_RETURN(uint64_t dim_u64, value.GetBitValueUnsigned());
             return GetArrayTypeColonAttr(array_type, dim_u64, node->attr());
           },

xls/dslx/constexpr_evaluator.cc

@@ -281,7 +281,9 @@ absl::Status ConstexprEvaluator::HandleColonRef(const ColonRef* expr) {
             XLS_ASSIGN_OR_RETURN(Expr * member_value_expr,
                                  enum_def->GetValue(expr->attr()));
 
-            // Since enum defs can't [currently] be parameterized, this is safe.
+            // Since enum defs can't [currently] be parameterized or declared
+            // locally, it's safe to assume we want to grab the root type info
+            // for the enum def.
             XLS_ASSIGN_OR_RETURN(
                 TypeInfo * type_info,
                 import_data_->GetRootTypeInfoForNode(enum_def));
@@ -302,14 +304,13 @@ absl::Status ConstexprEvaluator::HandleColonRef(const ColonRef* expr) {
             return absl::OkStatus();
           },
           [&](ArrayTypeAnnotation* array_type_annotation) -> absl::Status {
-            XLS_ASSIGN_OR_RETURN(
-                TypeInfo * type_info,
-                import_data_->GetRootTypeInfoForNode(array_type_annotation));
+            const TypeInfo& type_info = GetTypeInfoForNodeIfDifferentModule(
+                array_type_annotation, *type_info_, *import_data_);
             XLS_RET_CHECK(
-                type_info->IsKnownConstExpr(array_type_annotation->dim()));
+                type_info.IsKnownConstExpr(array_type_annotation->dim()));
             XLS_ASSIGN_OR_RETURN(
                 InterpValue dim,
-                type_info->GetConstExpr(array_type_annotation->dim()));
+                type_info.GetConstExpr(array_type_annotation->dim()));
             XLS_ASSIGN_OR_RETURN(uint64_t dim_u64, dim.GetBitValueViaSign());
             XLS_ASSIGN_OR_RETURN(InterpValue value,
                                  GetArrayTypeColonAttr(array_type_annotation,

xls/dslx/tests/BUILD

@@ -265,6 +265,10 @@ dslx_lang_test(
     dslx_deps = [":constexpr_dslx"],
 )
 
+dslx_lang_test(name = "attr_via_local_type_alias")
+
+dslx_lang_test(name = "attr_via_local_type_alias_in_parametric")
+
 dslx_lang_test(
     name = "builtin_type_max",
     dslx_deps = [":number_of_imported_type_import_dslx"],

xls/dslx/tests/attr_via_local_type_alias.x

@@ -0,0 +1,30 @@
+// Copyright 2025 The XLS Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// This sample shows use of a global constant definition in a local type alias.
+
+const N = u32:7;
+
+fn main() -> (uN[N], uN[N], uN[N]) {
+    type T = uN[N];
+    (T::MIN, T::ZERO, T::MAX)
+}
+
+#[test]
+fn test_f() {
+    let (min, zero, max) = main();
+    assert_eq(min, u7:0);
+    assert_eq(zero, u7:0);
+    assert_eq(max, u7:127);
+}

xls/dslx/tests/attr_via_local_type_alias_in_parametric.x

@@ -0,0 +1,26 @@
+// Copyright 2025 The XLS Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// This sample shows use of a parametric binding in a local type alias.
+
+fn f<N: u32>() -> uN[N] {
+    type UN = uN[N];
+    let max = UN::MAX;
+    max
+}
+
+fn main() -> u8 { f<u32:8>() }
+
+#[test]
+fn test_f() { assert_eq(main(), u8:255); }

xls/dslx/type_system/deduce_utils.cc

@@ -862,4 +862,18 @@ bool IsAcceptableCast(const Type& from, const Type& to) {
   return false;
 }
 
+const TypeInfo& GetTypeInfoForNodeIfDifferentModule(
+    AstNode* node, const TypeInfo& current_type_info,
+    const ImportData& import_data) {
+  if (node->owner() == current_type_info.module()) {
+    return current_type_info;
+  }
+  absl::StatusOr<const TypeInfo*> type_info =
+      import_data.GetRootTypeInfoForNode(node);
+  CHECK_OK(type_info.status())
+      << "Must be able to get root type info for node " << node->ToString();
+  CHECK(type_info.value() != nullptr);
+  return *type_info.value();
+}
+
 }  // namespace xls::dslx

xls/dslx/type_system/deduce_utils.h

@@ -218,6 +218,13 @@ absl::StatusOr<std::optional<Function*>> ImplFnFromCallee(
 // should not cause a type error to occur).
 bool IsAcceptableCast(const Type& from, const Type& to);
 
+// Returns the TypeInfo for the given node, preferring the current TypeInfo if
+// the node is in the same module, otherwise giving the root TypeInfo for
+// the node's module.
+const TypeInfo& GetTypeInfoForNodeIfDifferentModule(
+    AstNode* node, const TypeInfo& current_type_info,
+    const ImportData& import_data);
+
 }  // namespace xls::dslx
 
 #endif  // XLS_DSLX_TYPE_SYSTEM_DEDUCE_UTILS_H_

xls/dslx/type_system/typecheck_module_test.cc

@@ -4216,6 +4216,25 @@ fn main() { x36() }
   XLS_EXPECT_OK(Typecheck(kProgram));
 }
 
+// Previously this would cause us to RET_CHECK because we were assuming we
+// wanted to grab the root type information instead of the parametric
+// invocation's type information.
+TEST(TypecheckTest, AttrViaParametricBinding) {
+  constexpr std::string_view kProgram = R"(
+fn f<N: u32>() -> uN[N]{
+    type UN = uN[N];
+    let max = UN::MAX;
+    max
+}
+
+#[test]
+fn test_f() {
+    assert_eq(f<u32:8>(), u8:255);
+}
+)";
+  XLS_EXPECT_OK(Typecheck(kProgram));
+}
+
 // Table-oriented test that lets us validate that *types on parameters* are
 // compatible with *particular values* that should be type-compatible.
 TEST(PassValueToIdentityFnTest, ParameterVsValue) {