consolidation · T2L · Jan 13, 2020 · Jan 13, 2020 · Jan 13, 2020 · Jan 20, 2020
diff --git a/src/Comments.php b/src/Comments.php
@@ -27,6 +27,7 @@ class Comments
     protected $hasStored;
     protected $headComments;
     protected $accumulated;
+    protected $lineIds;
     protected $stored;
     protected $endComments;
 
@@ -35,6 +36,7 @@ public function __construct()
         $this->hasStored = false;
         $this->headComments = false;
         $this->accumulated = [];
+        $this->lineIds = [];
         $this->stored = [];
         $this->endComments = [];
     }
@@ -137,6 +139,7 @@ protected function accumulate($line)
      */
     protected function storeAccumulated($line)
     {
+
         // Remember that we called storeAccumulated at least once
         $this->hasStored = true;
 
@@ -150,26 +153,81 @@ protected function storeAccumulated($line)
             return;
         }
         if (!empty($this->accumulated)) {
-            // $this->stored saves a stack of accumulated results.
-            $this->stored[$line][] = $this->accumulated;
+            $lineId = $this->getLineId($line, true);
+            $this->stored[$lineId][] = $this->accumulated;
             $this->accumulated = [];
         }
     }
 
+    /**
+     * Generates unique line id based on the key it contains. It allows
+     * to reattach comments to the edited yaml sections.
+     *
+     * For example, lets take a look at the following yaml:
+     *
+     * # Top comments
+     * top:
+     * # Top one
+     *   one:
+     *     # Top two
+     *     two: two
+     * # Bottom comments
+     * bottom:
+     *   # Bottom one
+     *   one:
+     *     # Bottom two
+     *     two: 2
+     *
+     * This method generates ids based on keys (discarding values).
+     * Additionally, duplicating keys are taken into account as well.
+     * The following ids will be generated:
+     *
+     * 1|top
+     * 1|  one
+     * 1|    two
+     * 1|bottom
+     * 2|  one
+     * 2|    two
+     *
+     * @param string $line
+     * @param bool $isCollecting
+     */
+    protected function getLineId($line, $isCollecting = true)
+    {
+        list($id) = explode(':', $line, 2);
+
+        if ($isCollecting) {
+            if (isset($this->lineIds[$id])) {
+                $this->lineIds[$id][] = end($this->lineIds[$id]) + 1;
+            } else {
+                $this->lineIds[$id] = [1];
+            }
+
+            return end($this->lineIds[$id]) . '|' . $id;
+        }
+
+        if (isset($this->lineIds[$id])) {
+            return array_shift($this->lineIds[$id]) . '|' . $id;
+        } else {
+            return  '1|' . $id;
+        }
+    }
+
     /**
      * Check to see if the provided line has any associated comments.
      *
      * @param string $line
      */
     protected function find($line)
     {
-        if (!isset($this->stored[$line]) || empty($this->stored[$line])) {
+        $lineId = $this->getLineId($line, false);
+        if (!isset($this->stored[$lineId]) || empty($this->stored[$lineId])) {
             return [];
         }
         // The stored result is a stack of accumulated comments. Pop
         // one off; if more remain, they will be attached to the next
         // line with the same value.
-        return array_shift($this->stored[$line]);
+        return array_shift($this->stored[$lineId]);
     }
 
     /**

diff --git a/tests/TestComments.php b/tests/TestComments.php
@@ -1,6 +1,7 @@
 <?php
 namespace Consolidation\Comments;
 
+use Symfony\Component\Yaml\Parser;
 use PHPUnit\Framework\TestCase;
 
 class TestComments extends TestCase
@@ -51,6 +52,42 @@ function commentTestData()
     two:
 EOT;
 
+        $duplicate_comments = <<< 'EOT'
+# Top comments
+top:
+  # Top one
+  one:
+    # Top two
+    two:
+# Bottom comments
+bottom:
+  # Bottom one
+  one:
+    # Bottom two
+    two: two
+EOT;
+
+        $duplicate_altered_without_comments = <<< 'EOT'
+top:
+  one:
+    two: 2
+bottom:
+  one: 1
+EOT;
+
+        $duplicate_altered_with_comments = <<< 'EOT'
+# Top comments
+top:
+  # Top one
+  one:
+    # Top two
+    two: 2
+# Bottom comments
+bottom:
+  # Bottom one
+  one: 1
+EOT;
+
         $travis_yaml_with_comments = <<< 'EOT'
 dist: trusty
 language: php
@@ -128,6 +165,7 @@ function commentTestData()
         return [
             [ $simple_yaml, $simple_yaml_reordered, $simple_yaml_expected, ],
             [ $indented_comments, $indented_without_comments, $indented_comments, ],
+            [ $duplicate_comments, $duplicate_altered_without_comments, $duplicate_altered_with_comments, ],
             [ $travis_yaml_with_comments, $travis_yaml_without_comments, $travis_yaml_with_comments ],
         ];
     }
@@ -139,6 +177,13 @@ function commentTestData()
      */
     function testCommentParsing($original_contents, $altered_contents, $expected)
     {
+        // Ensure that passed args are valid YAML.
+
+        $parser = new Parser();
+        foreach (func_get_args() as $arg) {
+          $parser->parse($arg);
+        }
+
         // Second step: collect comments from original document and inject them into result.
 
         $commentManager = new Comments();