Merge pull request google#1842 from xlsynth:cdleary/2024-12-15-use-typecheck-interp

PiperOrigin-RevId: 716833766

Author: copybara-github

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_; }