Merge pull request #1842 from xlsynth:cdleary/2024-12-15-use-typechec…

…k-interp

PiperOrigin-RevId: 716833766

xls/dslx/BUILD

@@ -39,6 +39,8 @@ cc_library(
     hdrs = ["virtualizable_file_system.h"],
     deps = [
         "//xls/common/file:filesystem",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:str_format",

xls/dslx/bytecode/bytecode.h

@@ -51,7 +51,8 @@ class Bytecode {
     kAnd,
     // Invokes the function given in the Bytecode's data argument. Arguments are
     // given on the stack with deeper elements being earlier in the arg list
-    // (rightmost arg is TOS0 because we evaluate args left-to-right).
+    // (rightmost arg is TOS1 because we evaluate args left-to-right, TOS0 is
+    // the callee).
     kCall,
     // Casts the element on top of the stack to the type given in the optional
     // arg.

xls/dslx/bytecode/bytecode_emitter.cc

@@ -1329,6 +1329,49 @@ absl::Status BytecodeEmitter::HandleNameRef(const NameRef* node) {
   return absl::OkStatus();
 }
 
+absl::StatusOr<InterpValue> BytecodeEmitter::HandleExternRef(
+    const NameRef& name_ref, const NameDef& name_def,
+    UseTreeEntry& use_tree_entry) {
+  XLS_ASSIGN_OR_RETURN(const ImportedInfo* imported_info,
+                       type_info_->GetImportedOrError(&use_tree_entry));
+  Module* referenced_module = imported_info->module;
+  std::optional<ModuleMember*> member =
+      referenced_module->FindMemberWithName(name_def.identifier());
+  XLS_RET_CHECK(member.has_value());
+
+  // Note: we currently do not support re-exporting a `use` binding, when we do,
+  // this will need to potentially walk transitively to the definition.
+  std::optional<InterpValue> value = absl::visit(
+      Visitor{[&](Function* f) -> std::optional<InterpValue> {
+                // Type checking should have validated we refer to public
+                // members.
+                CHECK(f->is_public()) << f->ToString();
+                return InterpValue::MakeFunction(
+                    InterpValue::UserFnData{f->owner(), f});
+              },
+              [&](ConstantDef* cd) -> std::optional<InterpValue> {
+                // Type checking should have validated we refer to public
+                // members.
+                CHECK(cd->is_public()) << cd->ToString();
+                std::optional<InterpValue> value =
+                    imported_info->type_info->GetConstExprOption(cd->value());
+                // ConstantDef should always have a ConstExpr value associated
+                // with it.
+                CHECK(value.has_value());
+                return value;
+              },
+              [&](auto) -> std::optional<InterpValue> { return std::nullopt; }},
+      *member.value());
+  if (value.has_value()) {
+    return value.value();
+  }
+  return absl::InternalError(absl::StrFormat(
+      "Unhandled external reference to `%s` kind `%s` via `%s` @ %s",
+      name_def.identifier(), GetModuleMemberTypeName(*member.value()),
+      use_tree_entry.parent()->ToString(),
+      name_def.span().ToString(file_table())));
+}
+
 absl::StatusOr<std::variant<InterpValue, Bytecode::SlotIndex>>
 BytecodeEmitter::HandleNameRefInternal(const NameRef* node) {
   AnyNameDef any_name_def = node->name_def();
@@ -1350,6 +1393,11 @@ BytecodeEmitter::HandleNameRefInternal(const NameRef* node) {
                 CHECK(name_def != nullptr);
                 AstNode* definer = name_def->definer();
 
+                if (auto* use_tree_entry = dynamic_cast<UseTreeEntry*>(definer);
+                    use_tree_entry != nullptr) {
+                  return HandleExternRef(*node, *name_def, *use_tree_entry);
+                }
+
                 // Emit function and constant refs directly so that they can be
                 // stack elements without having to load slots with them.
                 if (auto* f = dynamic_cast<Function*>(definer); f != nullptr) {

xls/dslx/bytecode/bytecode_emitter.h

@@ -122,8 +122,14 @@ class BytecodeEmitter : public ExprVisitor {
   absl::Status HandleLet(const Let* node) override;
   absl::Status HandleMatch(const Match* node) override;
   absl::Status HandleNameRef(const NameRef* node) override;
+
   absl::StatusOr<std::variant<InterpValue, Bytecode::SlotIndex>>
   HandleNameRefInternal(const NameRef* node);
+
+  absl::StatusOr<InterpValue> HandleExternRef(const NameRef& name_ref,
+                                              const NameDef& name_def,
+                                              UseTreeEntry& use_tree_entry);
+
   absl::Status HandleNumber(const Number* node) override;
   struct FormattedInterpValue {
     InterpValue value;

xls/dslx/constexpr_evaluator.cc

@@ -442,7 +442,24 @@ absl::Status ConstexprEvaluator::HandleMatch(const Match* expr) {
   return InterpretExpr(expr);
 }
 
+absl::Status ConstexprEvaluator::HandleExternRef(const NameRef* name_ref,
+                                                 const NameDef* name_def,
+                                                 UseTreeEntry* use_tree_entry) {
+  const FileTable& file_table = *name_ref->owner()->file_table();
+  LOG(ERROR) << "HandleExternRef: " << name_ref->ToString() << " @ "
+             << name_ref->span().ToString(file_table);
+  return absl::OkStatus();
+}
+
 absl::Status ConstexprEvaluator::HandleNameRef(const NameRef* expr) {
+  if (std::holds_alternative<const NameDef*>(expr->name_def())) {
+    auto* name_def = std::get<const NameDef*>(expr->name_def());
+    if (auto* use_tree_entry = dynamic_cast<UseTreeEntry*>(name_def->definer());
+        use_tree_entry != nullptr) {
+      return HandleExternRef(expr, name_def, use_tree_entry);
+    }
+  }
+
   AstNode* name_def = ToAstNode(expr->name_def());
 
   if (type_info_->IsKnownNonConstExpr(name_def) ||

xls/dslx/constexpr_evaluator.h

@@ -102,6 +102,9 @@ class ConstexprEvaluator : public xls::dslx::ExprVisitor {
   absl::Status HandleXlsTuple(const XlsTuple* expr) override;
 
  private:
+  absl::Status HandleExternRef(const NameRef* name_ref, const NameDef* name_def,
+                               UseTreeEntry* use_tree_entry);
+
   ConstexprEvaluator(ImportData* import_data, TypeInfo* type_info,
                      WarningCollector* warning_collector,
                      ParametricEnv bindings, const Type* type)

xls/dslx/create_import_data.cc

@@ -38,10 +38,14 @@ ImportData CreateImportData(
   return import_data;
 }
 
-ImportData CreateImportDataForTest() {
-  ImportData import_data(xls::kDefaultDslxStdlibPath,
-                         /*additional_search_paths=*/{}, kDefaultWarningsSet,
-                         std::make_unique<RealFilesystem>());
+ImportData CreateImportDataForTest(
+    std::unique_ptr<VirtualizableFilesystem> vfs) {
+  if (vfs == nullptr) {
+    vfs = std::make_unique<RealFilesystem>();
+  }
+  absl::Span<const std::filesystem::path> additional_search_paths = {};
+  ImportData import_data(xls::kDefaultDslxStdlibPath, additional_search_paths,
+                         kDefaultWarningsSet, std::move(vfs));
   import_data.SetBytecodeCache(std::make_unique<BytecodeCache>(&import_data));
   return import_data;
 }

xls/dslx/create_import_data.h

@@ -36,7 +36,8 @@ ImportData CreateImportData(
 
 // Creates an ImportData with reasonable defaults (standard path to the stdlib
 // and no additional search paths).
-ImportData CreateImportDataForTest();
+ImportData CreateImportDataForTest(
+    std::unique_ptr<VirtualizableFilesystem> vfs = nullptr);
 
 std::unique_ptr<ImportData> CreateImportDataPtrForTest();
 

xls/dslx/frontend/ast.cc

@@ -829,6 +829,27 @@ std::string UseInteriorEntry::ToString() const {
   return absl::StrCat(identifier_, "::", subtrees_str);
 }
 
+void UseTreeEntry::LinearizeToSubjects(std::vector<std::string>& prefix,
+                                       std::vector<UseSubject>& results) {
+  return absl::visit(
+      Visitor{
+          [&](const UseInteriorEntry& interior) {
+            prefix.push_back(std::string{interior.identifier()});
+            for (UseTreeEntry* subtree : interior.subtrees()) {
+              subtree->LinearizeToSubjects(prefix, results);
+            }
+            prefix.pop_back();
+          },
+          [&](NameDef* name_def) {
+            CHECK(name_def != nullptr);
+            UseSubject result(prefix, *name_def, *this);
+            result.mutable_identifiers().push_back(name_def->identifier());
+            results.push_back(std::move(result));
+          },
+      },
+      payload_);
+}
+
 std::string UseTreeEntry::ToString() const {
   return absl::visit(
       Visitor{
@@ -870,6 +891,16 @@ std::vector<std::string> UseTreeEntry::GetLeafIdentifiers() const {
   return result;
 }
 
+UseSubject::UseSubject(std::vector<std::string> identifiers, NameDef& name_def,
+                       UseTreeEntry& use_tree_entry)
+    : identifiers_(std::move(identifiers)),
+      name_def_(&name_def),
+      use_tree_entry_(&use_tree_entry) {}
+
+std::string UseSubject::ToErrorString() const {
+  return absl::StrCat("`", absl::StrJoin(identifiers_, "::"), "`");
+}
+
 absl::Status UseTreeEntry::Accept(AstNodeVisitor* v) const {
   return v->HandleUseTreeEntry(this);
 }

xls/dslx/frontend/ast.h

@@ -1360,6 +1360,41 @@ class UseInteriorEntry {
   std::vector<UseTreeEntry*> subtrees_;
 };
 
+// Represents a record of a "leaf" imported by a `use` statement into the module
+// scope -- note the `name_def`, this is the name definition that will be bound
+// in the module scope.
+//
+// A `use` like `use foo::{bar, baz}` will result in two `UseSubject`s.
+class UseSubject {
+ public:
+  UseSubject(std::vector<std::string> identifiers, NameDef& name_def,
+             UseTreeEntry& use_tree_entry);
+
+  absl::Span<std::string const> identifiers() const { return identifiers_; }
+  const NameDef& name_def() const { return *name_def_; }
+  const UseTreeEntry& use_tree_entry() const { return *use_tree_entry_; }
+  UseTreeEntry& use_tree_entry() { return *use_tree_entry_; }
+
+  std::vector<std::string>& mutable_identifiers() { return identifiers_; }
+
+  // Returns a string that represents the subject of the `use` statement in
+  // the form of a colon-ref; e.g. `foo::bar::baz`.
+  //
+  // Note that the returned value is surrounded in backticks.
+  std::string ToErrorString() const;
+
+ private:
+  // The identifier in the subject path; e.g. in `use foo::bar::baz` the
+  // identifiers are {`foo`, `bar`, `baz`}.
+  std::vector<std::string> identifiers_;
+
+  // The name definition that will be bound in the module scope.
+  const NameDef* name_def_;
+
+  // Use tree entry that wraps the `name_def`.
+  UseTreeEntry* use_tree_entry_;
+};
+
 // Arbitrary entry (interior or leaf) in the `use` construct tree.
 class UseTreeEntry : public AstNode {
  public:
@@ -1385,6 +1420,11 @@ class UseTreeEntry : public AstNode {
   }
   const Span& span() const { return span_; }
 
+  // Note: this is non-const because we capture a mutable AST node pointer in
+  // the results.
+  void LinearizeToSubjects(std::vector<std::string>& prefix,
+                           std::vector<UseSubject>& results);
+
  private:
   std::variant<UseInteriorEntry, NameDef*> payload_;
   Span span_;
@@ -1426,13 +1466,27 @@ class Use : public AstNode {
 
   std::vector<AstNode*> GetChildren(bool want_types) const override;
 
+  // Returns all the identifiers of the `NameDef`s at the leaves of the tree.
   std::vector<std::string> GetLeafIdentifiers() const {
     return root_->GetLeafIdentifiers();
   }
   std::vector<NameDef*> GetLeafNameDefs() const {
     return root_->GetLeafNameDefs();
   }
 
+  // Returns a vector of `UseSubject`s that represent the "subjects" (i.e.
+  // individual imported entities that get a name binding in the module) of
+  // the `use` statement.
+  //
+  // Note: this is non-const because we capture a mutable AST node pointer in
+  // the results.
+  std::vector<UseSubject> LinearizeToSubjects() {
+    std::vector<std::string> prefix;
+    std::vector<UseSubject> results;
+    root_->LinearizeToSubjects(prefix, results);
+    return results;
+  }
+
   UseTreeEntry& root() { return *root_; }
   const UseTreeEntry& root() const { return *root_; }
   const Span& span() const { return span_; }

xls/dslx/frontend/ast_test.cc

@@ -362,5 +362,6 @@ TEST(AstTest, IsConstantEmptyArray) {
 
   EXPECT_TRUE(IsConstant(array));
 }
+
 }  // namespace
 }  // namespace xls::dslx

xls/dslx/frontend/ast_utils.cc

@@ -281,6 +281,18 @@ absl::StatusOr<InterpValue> GetArrayTypeColonAttr(
                                 [&] { return array_type->ToString(); });
 }
 
+std::optional<const UseTreeEntry*> IsExternNameRef(const NameRef& name_ref) {
+  const AstNode* definer = name_ref.GetDefiner();
+  if (definer == nullptr) {
+    return std::nullopt;
+  }
+  auto* use_tree_entry = dynamic_cast<const UseTreeEntry*>(definer);
+  if (use_tree_entry == nullptr) {
+    return std::nullopt;
+  }
+  return use_tree_entry;
+}
+
 // Attempts to evaluate an expression as a literal boolean.
 //
 // This has a few simple forms:

xls/dslx/frontend/ast_utils.h

@@ -84,6 +84,9 @@ absl::StatusOr<InterpValue> GetArrayTypeColonAttr(
     const ArrayTypeAnnotation* type, uint64_t constexpr_dim,
     std::string_view attr);
 
+// Returns a non-nullopt value if `name_ref` is bound by a `use` statement.
+std::optional<const UseTreeEntry*> IsExternNameRef(const NameRef& name_ref);
+
 // -- Template Metaprogramming helpers for dealing with AST node variants
 
 // TMP helper that gets the Nth type from a parameter pack.

xls/dslx/frontend/module.h

@@ -172,7 +172,7 @@ class Module : public AstNode {
     }
 
     return absl::NotFoundError(
-        absl::StrFormat("No %s in module %s with name \"%s\"",
+        absl::StrFormat("No %s in module `%s` with name `%s`",
                         T::GetDebugTypeName(), name_, target_name));
   }
 

xls/dslx/frontend/parser.cc

@@ -1562,7 +1562,9 @@ absl::StatusOr<UseTreeEntry*> Parser::ParseUseTreeEntry(Bindings& bindings) {
     // subsequent level.
     XLS_ASSIGN_OR_RETURN(NameDef * name_def, TokenToNameDef(tok));
     bindings.Add(name_def->identifier(), name_def);
-    return module_->Make<UseTreeEntry>(name_def, tok.span());
+    auto* use_tree_entry = module_->Make<UseTreeEntry>(name_def, tok.span());
+    name_def->set_definer(use_tree_entry);
+    return use_tree_entry;
   }
 
   // If we've gotten here we know there's a next level, we're just looking to

xls/dslx/frontend/parser_test.cc

@@ -1666,12 +1666,22 @@ fn foo() -> bar::T<2>[3] {
 }
 
 TEST_F(ParserTest, UseOneItemFromModule) {
-  RoundTrip(R"(#![feature(use_syntax)]
+  std::unique_ptr<Module> m = RoundTrip(R"(#![feature(use_syntax)]
 
 use foo::BAR;
 fn main() -> u32 {
     BAR
 })");
+  std::optional<ModuleMember*> member = m->FindMemberWithName("BAR");
+  ASSERT_TRUE(member.has_value());
+  ASSERT_TRUE(std::holds_alternative<Use*>(*member.value()));
+  Use* use = std::get<Use*>(*member.value());
+  std::vector<UseSubject> subjects = use->LinearizeToSubjects();
+  ASSERT_EQ(subjects.size(), 1);
+  EXPECT_EQ(subjects.at(0).identifiers(),
+            std::vector<std::string>({"foo", "BAR"}));
+  EXPECT_EQ(subjects.at(0).name_def().identifier(), "BAR");
+  EXPECT_EQ(subjects.at(0).ToErrorString(), "`foo::BAR`");
 }
 
 TEST_F(ParserTest, UseTwoItemsFromModule) {

xls/dslx/import_data.h

@@ -70,6 +70,10 @@ class ModuleInfo {
 // Hashable (usable in a flat hash map).
 class ImportTokens {
  public:
+  static ImportTokens FromSpan(absl::Span<const std::string> identifiers) {
+    return ImportTokens(
+        std::vector<std::string>(identifiers.begin(), identifiers.end()));
+  }
   static absl::StatusOr<ImportTokens> FromString(std::string_view module_name);
 
   explicit ImportTokens(std::vector<std::string> pieces)
@@ -230,7 +234,8 @@ class ImportData {
       const std::filesystem::path&, absl::Span<const std::filesystem::path>,
       WarningKindSet);
 
-  friend ImportData CreateImportDataForTest();
+  friend ImportData CreateImportDataForTest(
+      std::unique_ptr<VirtualizableFilesystem> vfs);
   friend std::unique_ptr<ImportData> CreateImportDataPtrForTest();
 
   ImportData(std::filesystem::path stdlib_path,