Honor column heading above a group when aligning records

Align Record Columns now keeps a table lined up under a heading produced
by Add Column Headers: when the line directly above a record group is a
`--` comment with exactly one word per column, the records are aligned
to it and the heading is re-synced via the same path Add Column Headers
uses, so the two commands are mutually idempotent.

Other comment lines never interfere: a descriptive comment (any other
word count, or separated by a blank line) is never treated as a heading,
and comments interspersed within the group are left untouched. The new
matchHeadingForGroup helper is the exact-one-word-per-column gate.

Author: magnesj

vscode-extension/README.md

@@ -235,10 +235,14 @@ up at the decimal point position). Keyword headers, comment lines, the closing
 every data line above and below the comment is aligned against a single shared set
 of column widths.
 
-Comments are ignored when aligning: columns are positioned from the record data
-alone, and any comment lines (whether above or within the group) are left exactly
-as they are. A descriptive comment above a table is never mistaken for a column
-heading.
+Comments are ignored when aligning, with one exception: a **column heading**
+directly above the group — a `--` comment with exactly one word per column, as
+produced by [Add Column Headers](#add-column-headers) — is honoured. The records
+are aligned to it and the heading is kept in sync, so a table stays lined up
+under its heading across repeated alignments. Every other comment line is left
+exactly as it is: a descriptive comment above a table (any other word count, or
+one separated by a blank line) is never mistaken for a heading, and comments
+interspersed within the group are untouched.
 
 Before:
 ```

vscode-extension/src/extension.ts

@@ -16,6 +16,7 @@ import {
   parseUdqExpressionLine,
   formatUdqBlock,
   buildHeadingAndAlignedRecords,
+  matchHeadingForGroup,
   tokenColumnCount,
   toggleLineComments,
 } from './formatting';
@@ -824,14 +825,37 @@ function computeAlignEdits(
     // Extract just the record entries for formatting
     const records = entries.filter(e => e.record !== null).map(e => e.record as RecordLine);
 
-    // Columns are aligned from the record data alone. Comment lines — whether
-    // above the group or interspersed within it — are ignored for alignment
-    // and left untouched.
+    // Columns are aligned from the record data. A column heading directly above
+    // the group (a `--` comment with one word per column, as produced by "Add
+    // Column Headers") is honoured: the data is aligned to it and the heading is
+    // kept in sync. Any other comment line — a descriptive comment above the
+    // table, a heading not directly adjacent, or comments interspersed within
+    // the group — is ignored for alignment and left untouched.
     if (records.length > 1) {
       // Check whether the owning keyword is excluded before emitting any edits.
       const activeKw = findActiveKeyword(document, new vscode.Position(i, 0));
       if (!excludedKeywords.has((activeKw ?? '').toUpperCase())) {
-        const formatted = formatRecordGroup(records);
+        // The candidate heading is the line immediately above the first record.
+        const headingLineNum = i - 1;
+        const headingWords =
+          headingLineNum >= 0
+            ? matchHeadingForGroup(document.lineAt(headingLineNum).text, nCols)
+            : null;
+
+        let formatted: string[];
+        if (headingWords) {
+          const built = buildHeadingAndAlignedRecords(records, headingWords);
+          formatted = built.formattedRecords;
+          const headingOrig = document.lineAt(headingLineNum).text;
+          if (built.heading !== headingOrig) {
+            edits.push(
+              vscode.TextEdit.replace(document.lineAt(headingLineNum).range, built.heading),
+            );
+          }
+        } else {
+          formatted = formatRecordGroup(records);
+        }
+
         let recordIdx = 0;
         for (const entry of entries) {
           if (entry.record === null) { continue; } // comment line — leave as-is

vscode-extension/src/formatting.test.ts

@@ -9,6 +9,7 @@ import {
   formatUdqExpressionGroup,
   formatUdqBlock,
   parseHeadingPositions,
+  matchHeadingForGroup,
   formatRecordGroupWithHeading,
   buildHeadingAndAlignedRecords,
   tokenColumnCount,
@@ -666,6 +667,39 @@ describe('parseHeadingPositions', () => {
   });
 });
 
+// ---------------------------------------------------------------------------
+// matchHeadingForGroup — the discriminator that keeps ordinary comment lines
+// from being treated as column headings during alignment.
+// ---------------------------------------------------------------------------
+
+describe('matchHeadingForGroup', () => {
+  test('matches a comment with exactly one word per column', () => {
+    expect(matchHeadingForGroup('-- WELL GROUP I J', 4)).toEqual(['WELL', 'GROUP', 'I', 'J']);
+    expect(matchHeadingForGroup('--  Sw   Krw', 2)).toEqual(['Sw', 'Krw']);
+  });
+
+  test('rejects a descriptive comment whose word count differs from the columns', () => {
+    // A prose comment above the table must not be honoured as a heading.
+    expect(matchHeadingForGroup('-- multiplies PERMZ in the upper layers', 8)).toBeNull();
+    expect(matchHeadingForGroup('-- WELL GROUP I', 4)).toBeNull(); // too few
+    expect(matchHeadingForGroup('-- WELL GROUP I J K', 4)).toBeNull(); // too many
+  });
+
+  test('rejects non-comment lines and empty comments', () => {
+    expect(matchHeadingForGroup("'PERMZ' 0.2 1 1", 4)).toBeNull();
+    expect(matchHeadingForGroup('--', 2)).toBeNull();
+    expect(matchHeadingForGroup('--    ', 2)).toBeNull();
+  });
+
+  test('never matches a single-column group', () => {
+    expect(matchHeadingForGroup('-- PORO', 1)).toBeNull();
+  });
+
+  test('tolerates indentation and extra spacing between words', () => {
+    expect(matchHeadingForGroup('   --   A      B  ', 2)).toEqual(['A', 'B']);
+  });
+});
+
 // ---------------------------------------------------------------------------
 // formatRecordGroupWithHeading
 // ---------------------------------------------------------------------------

vscode-extension/src/formatting.ts

@@ -497,6 +497,26 @@ export function formatUdqBlock(lines: string[]): string[] {
   return lines.map((line, i) => (parsed[i] !== null ? formatted[idx++] : line));
 }
 
+/**
+ * Decide whether a comment line directly above a record group is that group's
+ * column heading (one label per column) rather than a free-form descriptive
+ * comment. A heading must be a `--` comment whose word count equals the group's
+ * column count `nCols` (and there must be at least two columns).
+ *
+ * The exact-count rule is what keeps ordinary comment lines from being treated
+ * as headings: a prose comment almost never has exactly one word per column,
+ * and the caller only ever offers the line immediately above the table — a
+ * comment elsewhere, or one separated by a blank line, is never considered.
+ * Returns the heading words on a match, else null.
+ */
+export function matchHeadingForGroup(commentLine: string, nCols: number): string[] | null {
+  if (nCols < 2) return null;
+  const m = commentLine.match(/^\s*--\s*(\S.*)$/);
+  if (!m) return null;
+  const words = m[1].trim().split(/\s+/).filter(Boolean);
+  return words.length === nCols ? words : null;
+}
+
 // Parse absolute char positions of each word in a heading comment line (-- word1 word2 ...)
 export function parseHeadingPositions(line: string): number[] | null {
   const m = line.match(/^(\s*--\s*)(.*)/);