WIP: Improve constexpr-if string formatting

Signed-off-by: Robert Winkler <[email protected]>

Author: rw1nkler

xls/dslx/fmt/ast_fmt.cc

@@ -1597,14 +1597,19 @@ DocRef Fmt(const String& n, Comments& comments, DocArena& arena) {
 //  if <break1> $test_expr <break1> {
 DocRef MakeConditionalTest(const Conditional& n, Comments& comments,
                            DocArena& arena) {
-  return ConcatN(
-      arena, {
-                 arena.Make(Keyword::kIf),
-                 arena.space(),
-                 FmtExpr(*n.test(), comments, arena, /*suppress_parens=*/true),
-                 arena.space(),
-                 arena.ocurl(),
-             });
+  std::vector<DocRef> pieces;
+  if (n.IsConst() && !n.IsPartOfLadder()) {
+    pieces.push_back(arena.Make(Keyword::kConst));
+    pieces.push_back(arena.space());
+  }
+  pieces.push_back(arena.Make(Keyword::kIf));
+  pieces.push_back(arena.space());
+  pieces.push_back(
+      FmtExpr(*n.test(), comments, arena, /*suppress_parens=*/true));
+  pieces.push_back(arena.space());
+  pieces.push_back(arena.ocurl());
+
+  return ConcatNGroup(arena, pieces);
 }
 
 // When there's an else-if, or multiple statements inside of the blocks, we

xls/dslx/fmt/ast_fmt_test.cc

@@ -558,6 +558,33 @@ TEST_F(FunctionFmtTest, ConditionalWithUnnecessaryParens) {
   EXPECT_EQ(got, want);
 }
 
+TEST_F(FunctionFmtTest, ConstexprConditionalInTernaryStyle) {
+  const std::string_view original =
+      "fn f<TEST:bool>(y:u32,z:u32)->u32{const if TEST{y}else{z}}";
+  XLS_ASSERT_OK_AND_ASSIGN(std::string got, DoFmt(original));
+  const std::string_view want =
+      R"(fn f<TEST: bool>(y: u32, z: u32) -> u32 { const if TEST { y } else { z } })";
+  EXPECT_EQ(got, want);
+}
+
+TEST_F(FunctionFmtTest, ConstexprConditionalWithElseIf) {
+  const std::string_view original =
+      "fn f<TEST: bool[2]>(x:u32[3])->u32{const if TEST[0]{x[0]}else if "
+      "TEST[1]{x[1]}else{x[2]}}";
+  XLS_ASSERT_OK_AND_ASSIGN(std::string got, DoFmt(original));
+  const std::string_view want =
+      R"(fn f<TEST: bool[2]>(x: u32[3]) -> u32 {
+    const if TEST[0] {
+        x[0]
+    } else if TEST[1] {
+        x[1]
+    } else {
+        x[2]
+    }
+})";
+  EXPECT_EQ(got, want);
+}
+
 TEST_F(FunctionFmtTest, SimpleForOneStatementNoTypeAnnotation) {
   const std::string_view original =
       "fn f(x:u32)->u32{for(i,accum)in u32:0..u32:4{accum+i}(x)}";

xls/dslx/frontend/ast.cc

@@ -719,23 +719,21 @@ std::vector<StatementBlock*> Conditional::GatherBlocks() {
 }
 
 std::string Conditional::ToStringInternal() const {
-  const absl::string_view const_prefix = is_const_ ? "const " : "";
+  const absl::string_view const_prefix = (is_const_ && IsPartOfLadder()) ? "const " : "";
 
   auto make_string = [&](std::string (AstNode::*to_str_fn)()
                              const) -> std::string {
     std::string test_str = (test_->*to_str_fn)();
     std::string then_str = (consequent_->*to_str_fn)();
-
-    std::string else_part = " ";
+    std::string else_str = "";
     if (has_else_) {
-      std::string else_str = (ToAstNode(alternate_)->*to_str_fn)();
-      absl::string_view view = else_str;
-      absl::ConsumePrefix(&view, const_prefix);
-      else_part = absl::StrFormat(" else %s", view);
+      else_str = absl::StrFormat(" else %s",
+          (ToAstNode(alternate_)->*to_str_fn)()
+      );
     }
 
     return absl::StrFormat("%sif %s %s%s", const_prefix, test_str, then_str,
-                           else_part);
+                           else_str);
   };
 
   std::string inline_str = make_string(&AstNode::ToInlineString);

xls/dslx/frontend/ast.h

@@ -2349,6 +2349,19 @@ class Conditional : public Expr {
   // blocks of this if/else-if.../else ladder.
   std::vector<StatementBlock*> GatherBlocks();
 
+  bool IsPartOfLadder() const {
+    if (parent() == nullptr) {
+      return false;
+    }
+
+    if (parent()->kind() == AstNodeKind::kConditional) {
+      AstNode *parent_node = parent();
+      Conditional *parent = down_cast<Conditional*>(parent_node);
+      return ToAstNode(parent->alternate()) == this;
+    }
+    return false;
+  }
+
  private:
   std::string ToStringInternal() const final;
 

xls/dslx/frontend/parser.cc

@@ -1172,6 +1172,7 @@ absl::StatusOr<Conditional*> Parser::ParseConditionalNode(
   for (StatementBlock* block : outer_conditional->GatherBlocks()) {
     block->SetEnclosing(outer_conditional);
   }
+  outer_conditional->SetParentage();
   return outer_conditional;
 }
 

xls/dslx/frontend/parser_test.cc

@@ -2940,6 +2940,21 @@ TEST_F(ParserTest, TernaryConditional) {
                 {"really_long_identifier_so_that_this_is_too_many_chars"});
 }
 
+TEST_F(ParserTest, ConstexprTernaryConditional) {
+  Expr* e = RoundTripExpr("const if true { u32:42 } else { u32:24 }", {});
+
+  EXPECT_TRUE(down_cast<Conditional*>(e)->IsConst());
+  EXPECT_FALSE(down_cast<Conditional*>(e)->HasElseIf());
+  EXPECT_FALSE(down_cast<Conditional*>(e)->HasMultiStatementBlocks());
+
+  RoundTripExpr(R"(const if really_long_identifier_so_that_this_is_too_many_chars {
+    u32:42
+} else {
+    u32:24
+})",
+                {"really_long_identifier_so_that_this_is_too_many_chars"});
+}
+
 TEST_F(ParserTest, ConditionalInBinopChain) {
   RoundTrip("const X = if true { u32:42 } else { u32:24 } + u32:1;");
   RoundTrip("const X = u32:1 + if true { u32:42 } else { u32:24 };");