From 53f30d068f9bb9e7888ecbe4c34273f5c3223e27 Mon Sep 17 00:00:00 2001 From: Erin Moore Date: Wed, 9 Oct 2024 08:56:59 -0700 Subject: [PATCH] Support (non-parametric) impl in dslx. https://github.com/google/xls/issues/1277 PiperOrigin-RevId: 684055362 --- xls/dslx/bytecode/bytecode_emitter.cc | 10 ++++ xls/dslx/bytecode/bytecode_emitter_test.cc | 35 ++++++++++++ xls/dslx/constexpr_evaluator.cc | 18 +++++++ xls/dslx/constexpr_evaluator_test.cc | 60 +++++++++++++++++++++ xls/dslx/frontend/ast.cc | 8 +++ xls/dslx/frontend/ast.h | 2 + xls/dslx/frontend/module.cc | 2 + xls/dslx/ir_convert/function_converter.cc | 16 +++++- xls/dslx/lsp/find_definition.cc | 5 +- xls/dslx/tests/BUILD | 2 + xls/dslx/tests/impl.x | 30 +++++++++++ xls/dslx/type_system/deduce.cc | 2 +- xls/dslx/type_system/deduce_utils.cc | 14 +++-- xls/dslx/type_system/deduce_utils.h | 8 ++- xls/dslx/type_system/impl_typecheck_test.cc | 16 ++++-- 15 files changed, 208 insertions(+), 20 deletions(-) create mode 100644 xls/dslx/tests/impl.x diff --git a/xls/dslx/bytecode/bytecode_emitter.cc b/xls/dslx/bytecode/bytecode_emitter.cc index 2af0d77dbb..00034003b8 100644 --- a/xls/dslx/bytecode/bytecode_emitter.cc +++ b/xls/dslx/bytecode/bytecode_emitter.cc @@ -855,6 +855,16 @@ absl::StatusOr BytecodeEmitter::HandleColonRefInternal( [&](Module* module) -> absl::StatusOr { return HandleColonRefToValue(module, node); }, + [&](StructDef* struct_def) -> absl::StatusOr { + std::optional constant_def = + struct_def->GetImplConstant(node->attr()); + if (!constant_def.has_value()) { + return absl::NotFoundError(absl::StrFormat( + "No impl with constant '%s' defined for struct '%s'", + node->attr(), struct_def->identifier())); + } + return type_info_->GetConstExpr(constant_def.value()); + }, }, resolved_subject); } diff --git a/xls/dslx/bytecode/bytecode_emitter_test.cc b/xls/dslx/bytecode/bytecode_emitter_test.cc index 8fd84a27ce..4570cdfce7 100644 --- a/xls/dslx/bytecode/bytecode_emitter_test.cc +++ b/xls/dslx/bytecode/bytecode_emitter_test.cc @@ -808,6 +808,41 @@ fn imported_enum_ref() -> import_0::ImportedEnum { IsOkAndHolds(InterpValue::MakeSBits(4, 2))); } +TEST(BytecodeEmitterTest, StructImplConstant) { + constexpr std::string_view kBaseProgram = R"( +struct Empty {} + +impl Empty { + const MY_CONST = u4:7; +} + +#[test] +fn struct_const_ref() -> u4 { + Empty::MY_CONST +} +)"; + + auto import_data = CreateImportDataForTest(); + XLS_ASSERT_OK_AND_ASSIGN( + TypecheckedModule tm, + ParseAndTypecheck(kBaseProgram, "test.x", "test", &import_data)); + + XLS_ASSERT_OK_AND_ASSIGN(TestFunction * tf, + tm.module->GetTest("struct_const_ref")); + XLS_ASSERT_OK_AND_ASSIGN(std::unique_ptr bf, + BytecodeEmitter::Emit(&import_data, tm.type_info, + tf->fn(), ParametricEnv())); + + const std::vector& bytecodes = bf->bytecodes(); + ASSERT_EQ(bytecodes.size(), 1); + + const Bytecode* bc = bytecodes.data(); + ASSERT_EQ(bc->op(), Bytecode::Op::kLiteral); + ASSERT_TRUE(bc->has_data()); + EXPECT_THAT(bytecodes.at(0).value_data(), + IsOkAndHolds(InterpValue::MakeSBits(4, 7))); +} + TEST(BytecodeEmitterTest, ImportedConstant) { constexpr std::string_view kImportedProgram = R"(pub const MY_CONST = u3:2;)"; constexpr std::string_view kBaseProgram = R"( diff --git a/xls/dslx/constexpr_evaluator.cc b/xls/dslx/constexpr_evaluator.cc index fff1617b7d..cdce4ae276 100644 --- a/xls/dslx/constexpr_evaluator.cc +++ b/xls/dslx/constexpr_evaluator.cc @@ -389,6 +389,24 @@ absl::Status ConstexprEvaluator::HandleColonRef(const ColonRef* expr) { expr, type_info->GetConstExpr(constant_def->value()).value()); return absl::OkStatus(); }, + [&](StructDef* struct_def) -> absl::Status { + std::optional constant_def = + struct_def->GetImplConstant(expr->attr()); + if (!constant_def.has_value()) { + return absl::NotFoundError(absl::StrFormat( + "No impl with constant '%s' defined for struct '%s'", + expr->attr(), struct_def->identifier())); + } + XLS_RETURN_IF_ERROR(Evaluate(import_data_, type_info_, + warning_collector_, bindings_, + constant_def.value()->value())); + XLS_RET_CHECK( + type_info_->IsKnownConstExpr(constant_def.value()->value())); + type_info_->NoteConstExpr( + expr, type_info_->GetConstExpr(constant_def.value()->value()) + .value()); + return absl::OkStatus(); + }, }, subject); } diff --git a/xls/dslx/constexpr_evaluator_test.cc b/xls/dslx/constexpr_evaluator_test.cc index 4e1170147e..49bc0df79f 100644 --- a/xls/dslx/constexpr_evaluator_test.cc +++ b/xls/dslx/constexpr_evaluator_test.cc @@ -619,5 +619,65 @@ fn main() -> MyArray { EXPECT_THAT(value.GetLength(), IsOkAndHolds(4)); } +TEST(ConstexprEvaluatorTest, ImplWithConstantSimple) { + constexpr std::string_view kProgram = R"( +struct MyStruct {} + +impl MyStruct { + const STRUCT_CONST = u32:7; +} + +fn main() -> u32 { + MyStruct::STRUCT_CONST +} +)"; + + ImportData import_data(CreateImportDataForTest()); + XLS_ASSERT_OK_AND_ASSIGN( + TypecheckedModule tm, + ParseAndTypecheck(kProgram, "test.x", "test", &import_data)); + + XLS_ASSERT_OK_AND_ASSIGN(Function * f, + tm.module->GetMemberOrError("main")); + WarningCollector warnings(kAllWarningsSet); + XLS_ASSERT_OK(ConstexprEvaluator::Evaluate(&import_data, tm.type_info, + &warnings, ParametricEnv(), + f->body(), nullptr)); + XLS_ASSERT_OK_AND_ASSIGN(InterpValue value, + tm.type_info->GetConstExpr(f->body())); + EXPECT_EQ(value.GetBitValueViaSign().value(), 7); +} + +TEST(ConstexprEvaluatorTest, ImplWithConstantRefGlobal) { + constexpr std::string_view kProgram = R"( +const SIZE = u32:4; + +struct MyStruct {} + +impl MyStruct { + const STRUCT_CONST = u32:2 * SIZE; +} + +fn main() -> u32 { + MyStruct::STRUCT_CONST +} +)"; + + ImportData import_data(CreateImportDataForTest()); + XLS_ASSERT_OK_AND_ASSIGN( + TypecheckedModule tm, + ParseAndTypecheck(kProgram, "test.x", "test", &import_data)); + + XLS_ASSERT_OK_AND_ASSIGN(Function * f, + tm.module->GetMemberOrError("main")); + WarningCollector warnings(kAllWarningsSet); + XLS_ASSERT_OK(ConstexprEvaluator::Evaluate(&import_data, tm.type_info, + &warnings, ParametricEnv(), + f->body(), nullptr)); + XLS_ASSERT_OK_AND_ASSIGN(InterpValue value, + tm.type_info->GetConstExpr(f->body())); + EXPECT_EQ(value.GetBitValueViaSign().value(), 8); +} + } // namespace } // namespace xls::dslx diff --git a/xls/dslx/frontend/ast.cc b/xls/dslx/frontend/ast.cc index b59ab7a70d..0f0b4b806a 100644 --- a/xls/dslx/frontend/ast.cc +++ b/xls/dslx/frontend/ast.cc @@ -1296,6 +1296,14 @@ std::vector StructDef::GetMemberNames() const { return names; } +std::optional StructDef::GetImplConstant( + std::string_view constant_name) { + if (!impl_.has_value()) { + return std::nullopt; + } + return impl_.value()->GetConstant(constant_name); +} + // -- class Impl Impl::Impl(Module* owner, Span span, TypeAnnotation* struct_ref, diff --git a/xls/dslx/frontend/ast.h b/xls/dslx/frontend/ast.h index c3ab0bd1f0..13518dea30 100644 --- a/xls/dslx/frontend/ast.h +++ b/xls/dslx/frontend/ast.h @@ -2311,6 +2311,8 @@ class StructDef : public AstNode { std::optional impl() const { return impl_; } + std::optional GetImplConstant(std::string_view constant_name); + private: Span span_; NameDef* name_def_; diff --git a/xls/dslx/frontend/module.cc b/xls/dslx/frontend/module.cc index adab891dbb..007e8b8102 100644 --- a/xls/dslx/frontend/module.cc +++ b/xls/dslx/frontend/module.cc @@ -257,6 +257,8 @@ std::optional Module::FindMemberWithName( } } else if (std::holds_alternative(member)) { continue; // These have no name / binding. + } else if (std::holds_alternative(member)) { + continue; // These have no name / binding. } else { LOG(FATAL) << "Unhandled module member variant: " << ToAstNode(member)->GetNodeTypeName(); diff --git a/xls/dslx/ir_convert/function_converter.cc b/xls/dslx/ir_convert/function_converter.cc index b80c76520e..1f1ed54044 100644 --- a/xls/dslx/ir_convert/function_converter.cc +++ b/xls/dslx/ir_convert/function_converter.cc @@ -2801,7 +2801,21 @@ absl::Status FunctionConverter::HandleColonRef(const ColonRef* node) { DefConst(node, value); return absl::OkStatus(); }, - }, + [&](StructDef* struct_def) -> absl::Status { + std::optional constant_def = + struct_def->GetImplConstant(node->attr()); + if (!constant_def.has_value()) { + return absl::NotFoundError(absl::StrFormat( + "No impl with constant '%s' defined for struct '%s'", + node->attr(), struct_def->identifier())); + } + XLS_ASSIGN_OR_RETURN( + InterpValue iv, + current_type_info_->GetConstExpr(constant_def.value())); + XLS_ASSIGN_OR_RETURN(Value value, InterpValueToValue(iv)); + DefConst(node, value); + return absl::OkStatus(); + }}, subject); } diff --git a/xls/dslx/lsp/find_definition.cc b/xls/dslx/lsp/find_definition.cc index c0ce3355f4..76bfc877cf 100644 --- a/xls/dslx/lsp/find_definition.cc +++ b/xls/dslx/lsp/find_definition.cc @@ -35,8 +35,8 @@ namespace xls::dslx { namespace { const NameDef* GetNameDef( - const std::variant& colon_ref_subject, + const std::variant& colon_ref_subject, std::string_view attr) { return absl::visit( Visitor{ @@ -46,6 +46,7 @@ const NameDef* GetNameDef( return ModuleMemberGetNameDef(*member.value()); }, [&](EnumDef* e) -> const NameDef* { return e->GetNameDef(attr); }, + [&](StructDef* s) -> const NameDef* { return s->name_def(); }, [](BuiltinNameDef*) -> const NameDef* { return nullptr; }, [](ArrayTypeAnnotation*) -> const NameDef* { return nullptr; }, }, diff --git a/xls/dslx/tests/BUILD b/xls/dslx/tests/BUILD index a9bbce50ea..e1f4b4fef1 100644 --- a/xls/dslx/tests/BUILD +++ b/xls/dslx/tests/BUILD @@ -444,6 +444,8 @@ dslx_lang_test( dslx_lang_test(name = "struct_as_parametric") +dslx_lang_test(name = "impl") + dslx_lang_test(name = "subtract_to_negative") dslx_lang_test(name = "trace") diff --git a/xls/dslx/tests/impl.x b/xls/dslx/tests/impl.x new file mode 100644 index 0000000000..5fb33512f1 --- /dev/null +++ b/xls/dslx/tests/impl.x @@ -0,0 +1,30 @@ +// Copyright 2024 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. + +struct Point { x: u32, y: u32 } + +impl Point { + const MY_CONST = u32:5; +} + +fn main() -> u32 { + let p = Point { x: u32:7, y: u32:8 }; + p::MY_CONST +} + +#[test] +fn use_impl_const() { + let p = Point { x: u32:7, y: u32:8 }; + assert_eq(p::MY_CONST, u32:5); +} diff --git a/xls/dslx/type_system/deduce.cc b/xls/dslx/type_system/deduce.cc index 18a057172c..b1b602438c 100644 --- a/xls/dslx/type_system/deduce.cc +++ b/xls/dslx/type_system/deduce.cc @@ -1044,7 +1044,7 @@ static absl::StatusOr> DeduceColonRefToImpl( node->attr(), struct_def->identifier()), ctx->file_table()); } - return DeduceConstantDef(constant.value(), ctx); + return ctx->Deduce(constant.value()); } absl::StatusOr> DeduceColonRef(const ColonRef* node, diff --git a/xls/dslx/type_system/deduce_utils.cc b/xls/dslx/type_system/deduce_utils.cc index e20dcf3a7d..1451a946d7 100644 --- a/xls/dslx/type_system/deduce_utils.cc +++ b/xls/dslx/type_system/deduce_utils.cc @@ -437,25 +437,23 @@ absl::StatusOr ResolveColonRefSubjectForTypeChecking( td.value()); } -absl::StatusOr< - std::variant> +absl::StatusOr> ResolveColonRefSubjectAfterTypeChecking(ImportData* import_data, const TypeInfo* type_info, const ColonRef* colon_ref) { XLS_ASSIGN_OR_RETURN(auto result, ResolveColonRefSubjectForTypeChecking( import_data, type_info, colon_ref)); - using ReturnT = absl::StatusOr< - std::variant>; + using ReturnT = + absl::StatusOr>; return absl::visit( Visitor{ [](Module* x) -> ReturnT { return x; }, [](EnumDef* x) -> ReturnT { return x; }, [](BuiltinNameDef* x) -> ReturnT { return x; }, [](ArrayTypeAnnotation* x) -> ReturnT { return x; }, - [](StructDef*) -> ReturnT { - return absl::InternalError( - "After type checking colon-ref subject cannot be a StructDef"); - }, + [](StructDef* x) -> ReturnT { return x; }, [](ColonRef*) -> ReturnT { return absl::InternalError( "After type checking colon-ref subject cannot be a StructDef"); diff --git a/xls/dslx/type_system/deduce_utils.h b/xls/dslx/type_system/deduce_utils.h index 86f0543c13..6efe2781b6 100644 --- a/xls/dslx/type_system/deduce_utils.h +++ b/xls/dslx/type_system/deduce_utils.h @@ -66,9 +66,7 @@ absl::Status ValidateNumber(const Number& number, const Type& type); // * a module // * an enum definition // * a builtin type (with a constant item on it, a la `u7::MAX`) -// -// Struct definitions cannot currently have constant items on them, so this will -// have to be flagged by the type checker. +// * a constant defined via `impl` on a `StructDef`. absl::StatusOr> ResolveColonRefSubjectForTypeChecking(ImportData* import_data, @@ -78,8 +76,8 @@ ResolveColonRefSubjectForTypeChecking(ImportData* import_data, // Implementation of the above that can be called after type checking has been // performed, in which case we can eliminate some of the (invalid) possibilities // so they no longer need to be handled. -absl::StatusOr< - std::variant> +absl::StatusOr> ResolveColonRefSubjectAfterTypeChecking(ImportData* import_data, const TypeInfo* type_info, const ColonRef* colon_ref); diff --git a/xls/dslx/type_system/impl_typecheck_test.cc b/xls/dslx/type_system/impl_typecheck_test.cc index b5cfa675a1..b1f485dc8b 100644 --- a/xls/dslx/type_system/impl_typecheck_test.cc +++ b/xls/dslx/type_system/impl_typecheck_test.cc @@ -71,9 +71,7 @@ const GLOBAL_DIMS = NUM_DIMS; HasSubstr("Cannot find a definition"))); } -// TODO: https://github.com/google/xls/issues/1277 - Support assignment in -// expressions. -TEST(TypecheckTest, DISABLED_ImplConstantExtracted) { +TEST(TypecheckTest, ImplConstantExtracted) { constexpr std::string_view kProgram = R"( struct Point { x: u32, y: u32 } @@ -86,6 +84,18 @@ const GLOBAL_DIMS = Point::NUM_DIMS; XLS_EXPECT_OK(Typecheck(kProgram)); } +TEST(TypecheckErrorTest, ConstantExtractionWithoutImpl) { + constexpr std::string_view kProgram = R"( +struct Point { x: u32, y: u32 } + +const GLOBAL_DIMS = Point::NUM_DIMS; +)"; + EXPECT_THAT( + Typecheck(kProgram), + StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("Struct 'Point' has no impl defining 'NUM_DIMS'"))); +} + TEST(TypecheckErrorTest, ConstantAccessWithoutImplDef) { constexpr std::string_view kProgram = R"( struct Point { x: u32, y: u32 }