From 22f66303f5dc2dca50595a065d25b9d18ac45cfd Mon Sep 17 00:00:00 2001
From: David Fang <fangism@google.com>
Date: Mon, 6 Apr 2020 13:02:06 -0700
Subject: [PATCH] Handle case where first token of a partition is multi-line.

PiperOrigin-RevId: 305098932
---
 common/formatting/state_node.cc      |  8 ++++++--
 common/formatting/state_node_test.cc | 15 +++++++++++++++
 verilog/formatting/formatter_test.cc |  3 +++
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/common/formatting/state_node.cc b/common/formatting/state_node.cc
index ede1c7a59..9f0fbfcbb 100644
--- a/common/formatting/state_node.cc
+++ b/common/formatting/state_node.cc
@@ -144,8 +144,12 @@ int StateNode::_UpdateColumnPosition() {
       }
       // Penalize based on the column position that resulted in appending
       // text up to the first newline.
-      return prev_state->current_column +
-             current_format_token.before.spaces_required + first_newline_pos;
+      if (IsRootState()) {
+        return first_newline_pos;
+      } else {
+        return prev_state->current_column +
+               current_format_token.before.spaces_required + first_newline_pos;
+      }
     }
   }
 
diff --git a/common/formatting/state_node_test.cc b/common/formatting/state_node_test.cc
index 265a8fad5..553054647 100644
--- a/common/formatting/state_node_test.cc
+++ b/common/formatting/state_node_test.cc
@@ -658,6 +658,21 @@ TEST_F(StateNodeTestFixture, ConstructionAppendingPrevStateOverflow) {
   }
 }
 
+// Tests that column positions account for multiline tokens.
+TEST_F(StateNodeTestFixture, MultiLineTokenFront) {
+  const int kInitialIndent = 1;
+  const std::vector<TokenInfo> tokens = {{0, "a23456789\nb234"}};
+  Initialize(kInitialIndent, tokens);
+  auto& ftokens = pre_format_tokens_;
+  ftokens[0].before.spaces_required = 1;
+
+  // First token on line:
+  auto parent_state = std::make_shared<StateNode>(*uwline, style);
+  EXPECT_EQ(ABSL_DIE_IF_NULL(parent_state)->current_column,
+            4 /* length("b234") */);
+  EXPECT_EQ(parent_state->cumulative_cost, 0);
+}
+
 // Tests that newly calculated column positions account for multiline tokens.
 TEST_F(StateNodeTestFixture, MultiLineToken) {
   const int kInitialIndent = 1;
diff --git a/verilog/formatting/formatter_test.cc b/verilog/formatting/formatter_test.cc
index 751ce41b4..b14a14291 100644
--- a/verilog/formatting/formatter_test.cc
+++ b/verilog/formatting/formatter_test.cc
@@ -136,6 +136,9 @@ static const std::initializer_list<FormatterTestCase> kFormatterTestCases = {
     {"", ""},
     {"\n", "\n"},
     {"\n\n", "\n\n"},
+    {"\t//comment\n", "//comment\n"},
+    {"\t/*comment*/\n", "/*comment*/\n"},
+    {"\t/*multi-line\ncomment*/\n", "/*multi-line\ncomment*/\n"},
     // preprocessor test cases
     {"`include    \"path/to/file.vh\"\n", "`include \"path/to/file.vh\"\n"},
     {"`include    `\"path/to/file.vh`\"\n", "`include `\"path/to/file.vh`\"\n"},