diff --git a/xls/dslx/bytecode/bytecode_emitter.cc b/xls/dslx/bytecode/bytecode_emitter.cc index 9e592c920b..12867261b3 100644 --- a/xls/dslx/bytecode/bytecode_emitter.cc +++ b/xls/dslx/bytecode/bytecode_emitter.cc @@ -797,11 +797,10 @@ absl::StatusOr BytecodeEmitter::HandleColonRefInternal( return GetBuiltinNameDefColonAttr(builtin_name_def, node->attr()); }, [&](ArrayTypeAnnotation* array_type) -> absl::StatusOr { - 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()); }, diff --git a/xls/dslx/constexpr_evaluator.cc b/xls/dslx/constexpr_evaluator.cc index f440748a16..f7f215054a 100644 --- a/xls/dslx/constexpr_evaluator.cc +++ b/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, diff --git a/xls/dslx/tests/BUILD b/xls/dslx/tests/BUILD index 4f90e886b6..4a2c09f786 100644 --- a/xls/dslx/tests/BUILD +++ b/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"], diff --git a/xls/dslx/tests/attr_via_local_type_alias.x b/xls/dslx/tests/attr_via_local_type_alias.x new file mode 100644 index 0000000000..80d2e02e31 --- /dev/null +++ b/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); +} diff --git a/xls/dslx/tests/attr_via_local_type_alias_in_parametric.x b/xls/dslx/tests/attr_via_local_type_alias_in_parametric.x new file mode 100644 index 0000000000..70f14f129e --- /dev/null +++ b/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() -> uN[N] { + type UN = uN[N]; + let max = UN::MAX; + max +} + +fn main() -> u8 { f() } + +#[test] +fn test_f() { assert_eq(main(), u8:255); } diff --git a/xls/dslx/type_system/deduce_utils.cc b/xls/dslx/type_system/deduce_utils.cc index bc5b622a0c..96dad50593 100644 --- a/xls/dslx/type_system/deduce_utils.cc +++ b/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 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 diff --git a/xls/dslx/type_system/deduce_utils.h b/xls/dslx/type_system/deduce_utils.h index 2f6eac7ddc..47f289a265 100644 --- a/xls/dslx/type_system/deduce_utils.h +++ b/xls/dslx/type_system/deduce_utils.h @@ -218,6 +218,13 @@ absl::StatusOr> 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_ diff --git a/xls/dslx/type_system/typecheck_module_test.cc b/xls/dslx/type_system/typecheck_module_test.cc index cd3b7c4e91..32159f7d77 100644 --- a/xls/dslx/type_system/typecheck_module_test.cc +++ b/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() -> uN[N]{ + type UN = uN[N]; + let max = UN::MAX; + max +} + +#[test] +fn test_f() { + assert_eq(f(), 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) {