Preserve Inline comments

src/compose/compose-node.ts

@@ -95,8 +95,11 @@ export function composeNode(
   }
   if (spaceBefore) node.spaceBefore = true
   if (comment) {
-    if (token.type === 'scalar' && token.source === '') node.comment = comment
-    else node.commentBefore = comment
+    if (token.type === 'scalar' && token.source === '') {
+      node.comment = comment
+    } else {
+      node.commentBefore = comment
+    }
   }
   // @ts-expect-error Type checking misses meaning of isSrcToken
   if (ctx.options.keepSourceTokens && isSrcToken) node.srcToken = token
@@ -124,9 +127,12 @@ export function composeEmptyNode(
       onError(anchor, 'BAD_ALIAS', 'Anchor cannot be an empty string')
   }
   if (spaceBefore) node.spaceBefore = true
+  // Extend the range to include any trailing content (like comments) if it's beyond the current range end
+  if (end > node.range[2]) {
+    node.range[2] = end
+  }
   if (comment) {
     node.comment = comment
-    node.range[2] = end
   }
   return node
 }

src/compose/resolve-block-map.ts

@@ -1,7 +1,7 @@
 import type { ParsedNode } from '../nodes/Node.ts'
 import { Pair } from '../nodes/Pair.ts'
 import { YAMLMap } from '../nodes/YAMLMap.ts'
-import type { BlockMap } from '../parse/cst.ts'
+import type { BlockMap, SourceToken } from '../parse/cst.ts'
 import type { CollectionTag } from '../schema/types.ts'
 import type { ComposeContext, ComposeNode } from './compose-node.ts'
 import type { ComposeErrorHandler } from './composer.ts'
@@ -12,6 +12,49 @@ import { mapIncludes } from './util-map-includes.ts'
 
 const startColMsg = 'All mapping items must start at the same column'
 
+/**
+ * Extracts key comments from separator tokens.
+ * Key comments are single-line comments that appear immediately after the colon.
+ * Multi-line comments (where there are additional comments after a newline) should
+ * be left as value comments. Returns the extracted comment and modified separator tokens.
+ */
+function extractKeyComment(sep: SourceToken[]): { keyComment: string, modifiedSep: SourceToken[] } {
+  if (!sep.length) return { keyComment: '', modifiedSep: [] }
+
+  let mapValueIndIndex = -1
+  let keyComment = ''
+  let modifiedSep = [...sep]
+
+  // Find the map-value-ind (colon)
+  for (let i = 0; i < sep.length; i++) {
+    if (sep[i].type === 'map-value-ind') {
+      mapValueIndIndex = i
+      break
+    }
+  }
+  if (mapValueIndIndex === -1) return { keyComment: '', modifiedSep: sep }
+
+  // Only treat comments as key comments if they are immediately after the colon, with only spaces in between (no newline)
+  let i = mapValueIndIndex + 1
+  while (i < sep.length && sep[i].type === 'space') i++
+  if (i < sep.length && sep[i].type === 'comment') {
+    // Check for any newline between colon and comment
+    let hasNewline = false
+    for (let j = mapValueIndIndex + 1; j < i; j++) {
+      if (sep[j].type === 'newline') {
+        hasNewline = true
+        break
+      }
+    }
+    if (!hasNewline) {
+      // This comment is inline after the colon
+      keyComment = sep[i].source.substring(1) || ' '
+      modifiedSep = sep.slice(0, i).concat(sep.slice(i + 1))
+    }
+  }
+  return { keyComment, modifiedSep }
+}
+
 export function resolveBlockMap(
   { composeNode, composeEmptyNode }: ComposeNode,
   ctx: ComposeContext,
@@ -80,8 +123,17 @@ export function resolveBlockMap(
     if (mapIncludes(ctx, map.items, keyNode))
       onError(keyStart, 'DUPLICATE_KEY', 'Map keys must be unique')
 
+    // Extract key comments from separator tokens and modify tokens for value processing
+    const { keyComment, modifiedSep } = extractKeyComment(sep ?? [])
+
+    // Apply key comment to the key node if found
+    if (keyComment) {
+      if (keyNode.comment) keyNode.comment += '\n' + keyComment
+      else keyNode.comment = keyComment
+    }
+
     // value properties
-    const valueProps = resolveProps(sep ?? [], {
+    const valueProps = resolveProps(modifiedSep, {
       indicator: 'map-value-ind',
       next: value,
       offset: keyNode.range[2],

src/nodes/Node.ts

@@ -54,6 +54,9 @@ export abstract class NodeBase {
   /** A comment before this */
   declare commentBefore?: string | null
 
+  /** A comment after key: */
+  declare commentAfterKey?: string | null
+
   /**
    * The `[start, value-end, node-end]` character offsets for the part of the
    * source parsed into this node (undefined if not parsed). The `value-end`

src/parse/cst.ts

@@ -31,6 +31,11 @@ export interface SourceToken {
   offset: number
   indent: number
   source: string
+  commentAfterKey?: boolean
+  commentOnParentKey?: string
+  parentComment?: string
+  commentIsAfterKey?: boolean
+  context?: any
 }
 
 export interface ErrorToken {

src/parse/parser.ts

@@ -14,6 +14,8 @@ import type {
 import { prettyToken, tokenType } from './cst.ts'
 import { Lexer } from './lexer.ts'
 
+const DEBUG = false
+
 function includesToken(list: SourceToken[], type: SourceToken['type']) {
   for (let i = 0; i < list.length; ++i) if (list[i].type === type) return true
   return false
@@ -241,13 +243,81 @@ export class Parser {
     while (this.stack.length > 0) yield* this.pop()
   }
 
+  private getCurrentContext(parent: Token) {
+    const top = parent
+    if (!top) return { context: 'stream' }
+
+    switch (top.type) {
+      case 'block-map': {
+        const it = top.items[top.items.length - 1]
+        if (!it) return { context: 'map-start' }
+        if (it.value) return { context: 'map-value', key: it.key }
+        if (it.sep) return { context: 'map-separator', sep: it.sep }
+        return { context: 'map-key', key: it.key }
+      }
+      case 'block-seq': {
+        const it = top.items[top.items.length - 1]
+        if (!it) return { context: 'seq-start' }
+        if (it.value) return { context: 'seq-value' }
+        return { context: 'seq-item' }
+      }
+      case 'flow-collection': {
+        const it = top.items[top.items.length - 1]
+        if (!it) return { context: 'flow-start' }
+        if (it.value) return { context: 'flow-value', key: it.key }
+        if (it.sep) return { context: 'flow-separator' }
+        return { context: 'flow-key', key: it.key }
+      }
+      default:
+        return { context: 'other', parentType: top.type }
+    }
+  }
+
   private get sourceToken() {
     const st: SourceToken = {
       type: this.type as SourceToken['type'],
       offset: this.offset,
       indent: this.indent,
       source: this.source
     }
+
+    if (this.type === 'comment') {
+      const parent = this.peek(1)
+      const currentContext = this.getCurrentContext(parent)
+
+      if (DEBUG) {
+        st.context = currentContext
+      }
+
+      if (parent && 'items' in parent && parent.items && currentContext.context === 'map-separator') {
+        const it = parent.items[parent.items.length - 1]
+        if (it?.sep) {
+          if (DEBUG) st.parentComment = (it?.sep.find(st => st.type === 'comment') ?? {})?.source
+          // Check if this comment appears right after a map-value-ind token
+          const mapValueIndIndex = it.sep.findIndex(token => token.type === 'map-value-ind')
+          if (mapValueIndIndex !== -1) {
+            // Check if all tokens after map-value-ind are spaces (no newlines) followed by this comment
+            let allSpacesAfterMapValue = true
+            for (let i = mapValueIndIndex + 1; i < it.sep.length; i++) {
+              const token = it.sep[i]
+              if (token.type === 'newline') {
+                allSpacesAfterMapValue = false
+                break
+              } else if (token.type !== 'space') {
+                allSpacesAfterMapValue = false
+                break
+              }
+            }
+
+            // If all tokens after map-value-ind are spaces (no newlines), this comment is commentIsAfterKey
+            if (allSpacesAfterMapValue) {
+              st.commentIsAfterKey = true
+            }
+          }
+        }
+      }
+    }
+
     return st
   }
 

tests/doc/comments.ts

@@ -237,7 +237,6 @@ describe('parse comments', () => {
             {
               key: { commentBefore: 'c0', value: 'k1' },
               value: {
-                commentBefore: 'c1',
                 items: [{ value: 'v1' }, { commentBefore: 'c2', value: 'v2' }],
                 comment: 'c3'
               }
@@ -252,8 +251,7 @@ describe('parse comments', () => {
       })
       expect(String(doc)).toBe(source`
         #c0
-        k1:
-          #c1
+        k1: #c1
           - v1
           #c2
           - v2
@@ -1148,9 +1146,17 @@ entryB:
   })
 
   test('collection end comment', () => {
-    const src = `a: b #c\n#d\n`
+    const src = source`
+    a: b #c
+
+    #d
+    `
     const doc = YAML.parseDocument(src)
-    expect(String(doc)).toBe(`a: b #c\n\n#d\n`)
+    expect(String(doc)).toBe(source`
+    a: b #c
+
+    #d
+    `)
   })
 
   test('blank line after seq in map', () => {

tests/doc/stringify.ts

@@ -414,8 +414,7 @@ z:
       `
       const doc = YAML.parseDocument(src)
       expect(String(doc)).toBe(source`
-        key:
-          #comment
+        key: #comment
           - one
           - two
       `)
@@ -429,8 +428,7 @@ z:
       `
       const doc = YAML.parseDocument(src)
       expect(String(doc)).toBe(source`
-        key:
-          #comment
+        key: #comment
           !tag
           - one
           - two