fix(attach): end comments for the first item in block mapping/sequence

src/__snapshots__/attach.test.ts.snap

@@ -44,6 +44,56 @@ mappingValue (2:1 ~ 2:2)
 </mappingValue>
 `;
 
+exports[`"?\\n  # comment 1\\nb:\\n  # comment 2" 1`] = `
+comment (2:3 ~ 2:14)
+1 | ?¶
+2 | ··#·comment·1¶
+  |   ^^^^^^^^^^^
+3 | b:¶
+4 | ··#·comment·2¶
+mappingKey (1:1 ~ 2:14)
+<mappingKey>
+  <endComment value=" comment 1">
+</mappingKey>
+`;
+
+exports[`"?\\n  # comment 1\\nb:\\n  # comment 2" 2`] = `
+comment (4:3 ~ 4:14)
+2 | ··#·comment·1¶
+3 | b:¶
+4 | ··#·comment·2¶
+  |   ^^^^^^^^^^^
+mappingValue (3:2 ~ 4:14)
+<mappingValue>
+  <endComment value=" comment 2">
+</mappingValue>
+`;
+
+exports[`"?\\n  # comment 1\\nbb:\\n  # comment 2" 1`] = `
+comment (2:3 ~ 2:14)
+1 | ?¶
+2 | ··#·comment·1¶
+  |   ^^^^^^^^^^^
+3 | bb:¶
+4 | ··#·comment·2¶
+mappingKey (1:1 ~ 2:14)
+<mappingKey>
+  <endComment value=" comment 1">
+</mappingKey>
+`;
+
+exports[`"?\\n  # comment 1\\nbb:\\n  # comment 2" 2`] = `
+comment (4:3 ~ 4:14)
+2 | ··#·comment·1¶
+3 | bb:¶
+4 | ··#·comment·2¶
+  |   ^^^^^^^^^^^
+mappingValue (3:3 ~ 4:14)
+<mappingValue>
+  <endComment value=" comment 2">
+</mappingValue>
+`;
+
 exports[`"\\n# 123\\n%YAML 1.2\\n# 456\\n# 789\\n---\\n" 1`] = `
 comment (2:1 ~ 2:6)
 1 | ¶
@@ -1537,6 +1587,106 @@ mappingValue (2:7 ~ 4:11)
 </mappingValue>
 `;
 
+exports[`"- a\\n  # comment 1\\n- b\\n  # comment 2" 1`] = `
+comment (2:3 ~ 2:14)
+1 | -·a¶
+2 | ··#·comment·1¶
+  |   ^^^^^^^^^^^
+3 | -·b¶
+4 | ··#·comment·2¶
+sequenceItem (1:1 ~ 2:14)
+<sequenceItem>
+  <endComment value=" comment 1">
+</sequenceItem>
+`;
+
+exports[`"- a\\n  # comment 1\\n- b\\n  # comment 2" 2`] = `
+comment (4:3 ~ 4:14)
+2 | ··#·comment·1¶
+3 | -·b¶
+4 | ··#·comment·2¶
+  |   ^^^^^^^^^^^
+sequenceItem (3:1 ~ 4:14)
+<sequenceItem>
+  <endComment value=" comment 2">
+</sequenceItem>
+`;
+
+exports[`"- a\\n  # comment 1\\n- bb\\n  # comment 2" 1`] = `
+comment (2:3 ~ 2:14)
+1 | -·a¶
+2 | ··#·comment·1¶
+  |   ^^^^^^^^^^^
+3 | -·bb¶
+4 | ··#·comment·2¶
+sequenceItem (1:1 ~ 2:14)
+<sequenceItem>
+  <endComment value=" comment 1">
+</sequenceItem>
+`;
+
+exports[`"- a\\n  # comment 1\\n- bb\\n  # comment 2" 2`] = `
+comment (4:3 ~ 4:14)
+2 | ··#·comment·1¶
+3 | -·bb¶
+4 | ··#·comment·2¶
+  |   ^^^^^^^^^^^
+sequenceItem (3:1 ~ 4:14)
+<sequenceItem>
+  <endComment value=" comment 2">
+</sequenceItem>
+`;
+
+exports[`"a:\\n  # comment 1\\nb:\\n  # comment 2" 1`] = `
+comment (2:3 ~ 2:14)
+1 | a:¶
+2 | ··#·comment·1¶
+  |   ^^^^^^^^^^^
+3 | b:¶
+4 | ··#·comment·2¶
+mappingValue (1:2 ~ 2:14)
+<mappingValue>
+  <endComment value=" comment 1">
+</mappingValue>
+`;
+
+exports[`"a:\\n  # comment 1\\nb:\\n  # comment 2" 2`] = `
+comment (4:3 ~ 4:14)
+2 | ··#·comment·1¶
+3 | b:¶
+4 | ··#·comment·2¶
+  |   ^^^^^^^^^^^
+mappingValue (3:2 ~ 4:14)
+<mappingValue>
+  <endComment value=" comment 2">
+</mappingValue>
+`;
+
+exports[`"a:\\n  # comment 1\\nbb:\\n  # comment 2" 1`] = `
+comment (2:3 ~ 2:14)
+1 | a:¶
+2 | ··#·comment·1¶
+  |   ^^^^^^^^^^^
+3 | bb:¶
+4 | ··#·comment·2¶
+mappingValue (1:2 ~ 2:14)
+<mappingValue>
+  <endComment value=" comment 1">
+</mappingValue>
+`;
+
+exports[`"a:\\n  # comment 1\\nbb:\\n  # comment 2" 2`] = `
+comment (4:3 ~ 4:14)
+2 | ··#·comment·1¶
+3 | bb:¶
+4 | ··#·comment·2¶
+  |   ^^^^^^^^^^^
+mappingValue (3:3 ~ 4:14)
+<mappingValue>
+  <endComment value=" comment 2">
+</mappingValue>
+`;
+
 exports[`"a:\\n  [\\n    foo,\\n    # end comment\\n  ]" 1`] = `
 comment (4:5 ~ 4:18)
 2 | ··[¶

src/attach.test.ts

@@ -321,6 +321,30 @@ A:
   }`,
   `# --`,
   `# ---`,
+  `- a
+  # comment 1
+- b
+  # comment 2`,
+  `- a
+  # comment 1
+- bb
+  # comment 2`,
+  `a:
+  # comment 1
+b:
+  # comment 2`,
+  `a:
+  # comment 1
+bb:
+  # comment 2`,
+  `?
+  # comment 1
+b:
+  # comment 2`,
+  `?
+  # comment 1
+bb:
+  # comment 2`,
 ];
 
 testCases(cases.map((x): TestCaseSingle => [x, selector]), {

src/attach.ts

@@ -3,6 +3,7 @@ import {
   Document,
   EndCommentAttachable,
   LeadingCommentAttachable,
+  MappingKey,
   Root,
   TrailingCommentAttachable,
   YamlUnistNode,
@@ -162,6 +163,17 @@ function attachComment(
       currentNode = trailingNode;
     }
 
+    if (currentNode.type === "sequence" || currentNode.type === "mapping") {
+      currentNode = currentNode.children[0];
+    }
+
+    if (currentNode.type === "mappingItem") {
+      const [mappingKey, mappingValue] = currentNode.children;
+      currentNode = isExplicitMappingKey(mappingKey)
+        ? mappingKey
+        : mappingValue;
+    }
+
     while (true) {
       if (shouldOwnEndComment(currentNode, comment)) {
         defineParents(comment, currentNode);
@@ -226,13 +238,18 @@ function shouldOwnEndComment(
         (node.children.length === 0 ||
           (node.children.length === 1 &&
             node.children[0]!.type !== "blockFolded" &&
-            node.children[0]!.type !== "blockLiteral" &&
-            (node.type === "mappingValue" ||
-              // explicit key
-              node.position.start.offset !==
-                node.children[0]!.position.start.offset)))
+            node.children[0]!.type !== "blockLiteral")) &&
+        (node.type === "mappingValue" || isExplicitMappingKey(node))
       );
     default:
       return false;
   }
 }
+
+function isExplicitMappingKey(node: MappingKey) {
+  return (
+    node.position.start !== node.position.end &&
+    (node.children.length === 0 ||
+      node.position.start.offset !== node.children[0]!.position.start.offset)
+  );
+}

src/transform.ts

@@ -88,7 +88,7 @@ export function transformNode(
   node: YamlNode,
   context: Context,
 ): YamlUnistNode | null {
-  if (node === null) {
+  if (node === null || (node.type === undefined && node.value === null)) {
     return null;
   }