Add an 'operator' span to Binop and Unop.

This will unlock fixing #1739 PiperOrigin-RevId: 700439097
google · Nov 26, 2024 · 96a1ab0 · 96a1ab0
1 parent b3e17bb
commit 96a1ab0
Show file tree

Hide file tree

Showing 12 changed files with 96 additions and 71 deletions.
diff --git a/xls/codegen/vast/dslx_builder.cc b/xls/codegen/vast/dslx_builder.cc
@@ -439,13 +439,16 @@ absl::StatusOr<dslx::Expr*> DslxBuilder::ConvertMaxToWidth(
   dslx::Number* one = module().Make<dslx::Number>(
       span, "1", dslx::NumberKind::kOther, type_annot);
   return module().Make<dslx::Binop>(span, dslx::BinopKind::kAdd, casted_max,
-                                    one);
+                                    one, span);
 }
 
 dslx::Unop* DslxBuilder::HandleUnaryOperator(const dslx::Span& span,
                                              dslx::UnopKind unop_kind,
                                              dslx::Expr* arg) {
-  return module().Make<dslx::Unop>(span, dslx::UnopKind::kNegate, arg);
+  // Note it uses the same span for the whole node and the operand;
+  // the only time the operand span is used is for formatting, and
+  // this node won't be used for formatting.
+  return module().Make<dslx::Unop>(span, dslx::UnopKind::kNegate, arg, span);
 }
 
 absl::StatusOr<dslx::Expr*> DslxBuilder::HandleIntegerExponentiation(

diff --git a/xls/codegen/vast/translate_vast_to_dslx.cc b/xls/codegen/vast/translate_vast_to_dslx.cc
@@ -213,7 +213,10 @@ class VastToDslxTranslator {
             "%s vs %s",
             op->Emit(nullptr), lhs_type->ToString(), rhs_type->ToString()));
       }
-      result = module().Make<dslx::Binop>(span, kind, lhs, rhs);
+      // Note it uses the same span for the whole node and the operand;
+      // the only time the operand span is used is for formatting, and
+      // this node won't be used for formatting.
+      result = module().Make<dslx::Binop>(span, kind, lhs, rhs, span);
     }
     return dslx_builder_->MaybeCastToInferredVastType(op, result);
   }
@@ -266,11 +269,14 @@ class VastToDslxTranslator {
         dslx_expr =
             module().Make<dslx::Cast>(dslx_expr->span(), dslx_expr, annot);
       }
+      dslx::Span span = CreateNodeSpan(concat);
+      // Note it uses the same span for the whole node and the operand;
+      // the only time the operand span is used is for formatting, and
+      // this node won't be used for formatting.
       result = result == nullptr
                    ? dslx_expr
-                   : module().Make<dslx::Binop>(CreateNodeSpan(concat),
-                                                dslx::BinopKind::kConcat,
-                                                result, dslx_expr);
+                   : module().Make<dslx::Binop>(span, dslx::BinopKind::kConcat,
+                                                result, dslx_expr, span);
     }
     return result;
   }

diff --git a/xls/dslx/frontend/ast.cc b/xls/dslx/frontend/ast.cc
@@ -1612,9 +1612,10 @@ absl::StatusOr<BinopKind> BinopKindFromString(std::string_view s) {
 }
 
 Binop::Binop(Module* owner, Span span, BinopKind binop_kind, Expr* lhs,
-             Expr* rhs, bool in_parens)
+             Expr* rhs, Span op_span, bool in_parens)
     : Expr(owner, std::move(span), in_parens),
       binop_kind_(binop_kind),
+      op_span_(op_span),
       lhs_(lhs),
       rhs_(rhs) {}
 

diff --git a/xls/dslx/frontend/ast.h b/xls/dslx/frontend/ast.h
@@ -1388,9 +1388,10 @@ std::string UnopKindToString(UnopKind k);
 class Unop : public Expr {
  public:
   Unop(Module* owner, Span span, UnopKind unop_kind, Expr* operand,
-       bool in_parens = false)
+       Span op_span, bool in_parens = false)
       : Expr(owner, std::move(span), in_parens),
         unop_kind_(unop_kind),
+        op_span_(op_span),
         operand_(operand) {}
 
   ~Unop() override;
@@ -1412,6 +1413,8 @@ class Unop : public Expr {
 
   UnopKind unop_kind() const { return unop_kind_; }
   Expr* operand() const { return operand_; }
+  // The span of the operator, e.g., `++` for `x ++ y`.
+  Span op_span() const { return op_span_; }
 
   Precedence GetPrecedenceWithoutParens() const final {
     return Precedence::kUnaryOp;
@@ -1422,6 +1425,7 @@ class Unop : public Expr {
 
   UnopKind unop_kind_;
   Expr* operand_;
+  Span op_span_;
 };
 
 #define XLS_DSLX_BINOP_KIND_EACH(X)       \
@@ -1484,7 +1488,7 @@ const absl::btree_set<BinopKind>& GetBinopShifts();
 class Binop : public Expr {
  public:
   Binop(Module* owner, Span span, BinopKind kind, Expr* lhs, Expr* rhs,
-        bool in_parens = false);
+        Span op_span, bool in_parens = false);
 
   ~Binop() override;
 
@@ -1505,13 +1509,16 @@ class Binop : public Expr {
   BinopKind binop_kind() const { return binop_kind_; }
   Expr* lhs() const { return lhs_; }
   Expr* rhs() const { return rhs_; }
+  // The span of the operator, e.g., `++` for `x ++ y`.
+  Span op_span() const { return op_span_; }
 
   Precedence GetPrecedenceWithoutParens() const final;
 
  private:
   std::string ToStringInternal() const final;
 
   BinopKind binop_kind_;
+  Span op_span_;
   Expr* lhs_;
   Expr* rhs_;
 };

diff --git a/xls/dslx/frontend/ast_cloner.cc b/xls/dslx/frontend/ast_cloner.cc
@@ -99,7 +99,8 @@ class AstCloner : public AstNodeVisitor {
     XLS_RETURN_IF_ERROR(VisitChildren(n));
     old_to_new_[n] = module_->Make<Binop>(
         n->span(), n->binop_kind(), down_cast<Expr*>(old_to_new_.at(n->lhs())),
-        down_cast<Expr*>(old_to_new_.at(n->rhs())), n->in_parens());
+        down_cast<Expr*>(old_to_new_.at(n->rhs())), n->op_span(),
+        n->in_parens());
     return absl::OkStatus();
   }
 
@@ -926,9 +927,10 @@ class AstCloner : public AstNodeVisitor {
 
   absl::Status HandleUnop(const Unop* n) override {
     XLS_RETURN_IF_ERROR(VisitChildren(n));
-    old_to_new_[n] = module_->Make<Unop>(
-        n->span(), n->unop_kind(),
-        down_cast<Expr*>(old_to_new_.at(n->operand())), n->in_parens());
+    old_to_new_[n] =
+        module_->Make<Unop>(n->span(), n->unop_kind(),
+                            down_cast<Expr*>(old_to_new_.at(n->operand())),
+                            n->op_span(), n->in_parens());
     return absl::OkStatus();
   }
 

diff --git a/xls/dslx/frontend/ast_test.cc b/xls/dslx/frontend/ast_test.cc
@@ -262,7 +262,8 @@ TEST(AstTest, IsConstantBinopOfNameRefs) {
       m.Make<NameDef>(fake_span, std::string("MyStruct"), nullptr);
   NameRef* left = m.Make<NameRef>(fake_span, "name", name_def);
   NameRef* right = m.Make<NameRef>(fake_span, "name", name_def);
-  Binop* binop = m.Make<Binop>(fake_span, BinopKind::kAdd, left, right);
+  Binop* binop =
+      m.Make<Binop>(fake_span, BinopKind::kAdd, left, right, Span::Fake());
 
   EXPECT_FALSE(IsConstant(binop));
 }
@@ -276,7 +277,8 @@ TEST(AstTest, IsConstantBinopOfNumbers) {
                                 NumberKind::kOther, /*type=*/nullptr);
   Number* right = m.Make<Number>(fake_span, std::string("42"),
                                  NumberKind::kOther, /*type=*/nullptr);
-  Binop* binop = m.Make<Binop>(fake_span, BinopKind::kAdd, left, right);
+  Binop* binop =
+      m.Make<Binop>(fake_span, BinopKind::kAdd, left, right, Span::Fake());
 
   EXPECT_TRUE(IsConstant(binop));
 }
@@ -289,7 +291,7 @@ TEST(AstTest, IsConstantUnopOfNameRefs) {
   NameDef* name_def =
       m.Make<NameDef>(fake_span, std::string("MyStruct"), nullptr);
   NameRef* operand = m.Make<NameRef>(fake_span, "name", name_def);
-  Unop* unop = m.Make<Unop>(fake_span, UnopKind::kNegate, operand);
+  Unop* unop = m.Make<Unop>(fake_span, UnopKind::kNegate, operand, fake_span);
 
   EXPECT_FALSE(IsConstant(unop));
 }
@@ -301,7 +303,7 @@ TEST(AstTest, IsConstantUnopOfNumbers) {
 
   Number* operand = m.Make<Number>(fake_span, std::string("41"),
                                    NumberKind::kOther, /*type=*/nullptr);
-  Unop* unop = m.Make<Unop>(fake_span, UnopKind::kNegate, operand);
+  Unop* unop = m.Make<Unop>(fake_span, UnopKind::kNegate, operand, fake_span);
 
   EXPECT_TRUE(IsConstant(unop));
 }

diff --git a/xls/dslx/frontend/ast_test_utils.cc b/xls/dslx/frontend/ast_test_utils.cc
@@ -54,7 +54,7 @@ std::tuple<FileTable, Module, Binop*> MakeCastWithinLtComparison() {
 
   // x as t < x
   Cast* cast = m.Make<Cast>(fake_span, x_ref, type_ref_type_annotation);
-  Binop* lt = m.Make<Binop>(fake_span, BinopKind::kLt, cast, x_ref);
+  Binop* lt = m.Make<Binop>(fake_span, BinopKind::kLt, cast, x_ref, fake_span);
   return std::make_tuple(file_table, std::move(m), lt);
 }
 
@@ -134,7 +134,7 @@ std::tuple<FileTable, Module, Unop*> MakeCastWithinNegateExpression() {
 
   // x as u32
   Cast* cast = m.Make<Cast>(fake_span, x_ref, builtin_u32);
-  Unop* unop = m.Make<Unop>(fake_span, UnopKind::kNegate, cast);
+  Unop* unop = m.Make<Unop>(fake_span, UnopKind::kNegate, cast, fake_span);
   return std::make_tuple(file_table, std::move(m), unop);
 }
 
@@ -147,7 +147,8 @@ std::tuple<FileTable, Module, Attr*> MakeArithWithinAttrExpression() {
   BuiltinNameDef* y_def = m.GetOrCreateBuiltinNameDef("y");
   NameRef* y_ref = m.Make<NameRef>(fake_span, "y", y_def);
 
-  Binop* binop = m.Make<Binop>(fake_span, BinopKind::kMul, x_ref, y_ref);
+  Binop* binop =
+      m.Make<Binop>(fake_span, BinopKind::kMul, x_ref, y_ref, fake_span);
   Attr* attr = m.Make<Attr>(fake_span, binop, "my_attr");
   return std::make_tuple(file_table, std::move(m), attr);
 }

diff --git a/xls/dslx/frontend/parser.cc b/xls/dslx/frontend/parser.cc
@@ -1305,7 +1305,7 @@ absl::StatusOr<Expr*> Parser::ParseBinopChain(
       XLS_ASSIGN_OR_RETURN(BinopKind kind,
                            BinopKindFromString(TokenKindToString(op.kind())));
       Span span(lhs->span().start(), rhs->span().limit());
-      lhs = module_->Make<Binop>(span, kind, lhs, rhs);
+      lhs = module_->Make<Binop>(span, kind, lhs, rhs, op.span());
     } else {
       break;
     }
@@ -1413,7 +1413,7 @@ absl::StatusOr<Expr*> Parser::ParseComparisonExpression(
                               "comparison operators cannot be chained");
     }
     Span span(lhs->span().start(), rhs->span().limit());
-    lhs = module_->Make<Binop>(span, kind, lhs, rhs);
+    lhs = module_->Make<Binop>(span, kind, lhs, rhs, op.span());
   }
   VLOG(5) << "ParseComparisonExpression; result: `" << lhs->ToString() << "`";
   return lhs;
@@ -1781,7 +1781,7 @@ absl::StatusOr<Expr*> Parser::ParseTermLhs(Bindings& outer_bindings,
         LOG(FATAL) << "Inconsistent unary operation token kind.";
     }
     Span span(tok.span().start(), arg->span().limit());
-    lhs = module_->Make<Unop>(span, unop_kind, arg);
+    lhs = module_->Make<Unop>(span, unop_kind, arg, tok.span());
   } else if (peek->IsTypeKeyword() ||
              (peek->kind() == TokenKind::kIdentifier &&
               outer_bindings.ResolveNodeIsTypeDefinition(*peek->GetValue()))) {

diff --git a/xls/dslx/type_system/type_test.cc b/xls/dslx/type_system/type_test.cc
@@ -162,7 +162,7 @@ StructType CreateStructWithParametricExpr(Module& module) {
   NameDef* type_n_name_def = module.Make<NameDef>(
       Span::Fake(), "N",
       module.Make<Binop>(Span::Fake(), BinopKind::kAdd, type_m_name_ref,
-                         type_m_name_ref));
+                         type_m_name_ref, Span::Fake()));
   bindings.push_back(
       module.Make<ParametricBinding>(type_n_name_def, u32_type_annot, nullptr));
 

diff --git a/xls/dslx/type_system_v2/inference_table_test.cc b/xls/dslx/type_system_v2/inference_table_test.cc
@@ -159,8 +159,8 @@ TEST_F(InferenceTableTest, SignednessMismatch) {
   NameDef* y = module_->Make<NameDef>(Span::Fake(), "y", /*definer=*/nullptr);
   NameRef* x_ref = module_->Make<NameRef>(Span::Fake(), "x", x);
   NameRef* y_ref = module_->Make<NameRef>(Span::Fake(), "y", y);
-  AstNode* add_node =
-      module_->Make<Binop>(Span::Fake(), BinopKind::kAdd, x_ref, y_ref);
+  AstNode* add_node = module_->Make<Binop>(Span::Fake(), BinopKind::kAdd, x_ref,
+                                           y_ref, Span::Fake());
 
   TypeAnnotation* u32_annotation = module_->Make<BuiltinTypeAnnotation>(
       Span::Fake(), BuiltinType::kU32,
@@ -190,8 +190,8 @@ TEST_F(InferenceTableTest, SignednessAgreement) {
   NameDef* y = module_->Make<NameDef>(Span::Fake(), "y", /*definer=*/nullptr);
   NameRef* x_ref = module_->Make<NameRef>(Span::Fake(), "x", x);
   NameRef* y_ref = module_->Make<NameRef>(Span::Fake(), "y", y);
-  AstNode* add_node =
-      module_->Make<Binop>(Span::Fake(), BinopKind::kAdd, x_ref, y_ref);
+  AstNode* add_node = module_->Make<Binop>(Span::Fake(), BinopKind::kAdd, x_ref,
+                                           y_ref, Span::Fake());
 
   TypeAnnotation* u32_annotation = module_->Make<BuiltinTypeAnnotation>(
       Span::Fake(), BuiltinType::kU32,