From 2a6179cd03b5e19bac67ae17515f82db3b0e5537 Mon Sep 17 00:00:00 2001
From: Rowan Cockett <rowanc1@gmail.com>
Date: Fri, 13 Oct 2023 14:19:53 -0600
Subject: [PATCH] =?UTF-8?q?=F0=9F=AA=86=20Improve=20block=20nesting=20tran?=
 =?UTF-8?q?sform?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .changeset/five-cats-punch.md               |  5 +++
 packages/myst-transforms/src/blocks.spec.ts | 35 ++++++++++++++++++++-
 packages/myst-transforms/src/blocks.ts      | 35 ++++++++++++---------
 3 files changed, 60 insertions(+), 15 deletions(-)
 create mode 100644 .changeset/five-cats-punch.md

diff --git a/.changeset/five-cats-punch.md b/.changeset/five-cats-punch.md
new file mode 100644
index 000000000..d6b20bee1
--- /dev/null
+++ b/.changeset/five-cats-punch.md
@@ -0,0 +1,5 @@
+---
+'myst-transforms': patch
+---
+
+Block nesting updates
diff --git a/packages/myst-transforms/src/blocks.spec.ts b/packages/myst-transforms/src/blocks.spec.ts
index 30c556fb0..5436aca77 100644
--- a/packages/myst-transforms/src/blocks.spec.ts
+++ b/packages/myst-transforms/src/blocks.spec.ts
@@ -1,9 +1,42 @@
 import { describe, expect, test } from 'vitest';
 import { u } from 'unist-builder';
 import { VFile } from 'vfile';
-import { blockMetadataTransform } from './blocks';
+import { blockMetadataTransform, blockNestingTransform } from './blocks';
 
 describe('Test blockMetadataTransform', () => {
+  test.each([
+    [
+      ['a', 'block', 'b'],
+      ['block>a', 'block', 'block>b'],
+    ],
+    [
+      ['block', 'b'],
+      ['block', 'block>b'],
+    ],
+    [['a', 'b'], ['block>a,b']],
+    [
+      ['block', 'block'],
+      ['block', 'block'],
+    ],
+    [
+      ['a', 'block', 'b', 'c'],
+      ['block>a', 'block', 'block>b,c'],
+    ],
+  ])('nestBlock(%s, %s)}', (a, b) => {
+    const mdast = u(
+      'root',
+      a.map((type) => ({ type })),
+    ) as any;
+    blockNestingTransform(mdast);
+    expect(
+      mdast.children.map(
+        ({ type, children }) =>
+          `${type}${
+            children && children.length > 0 ? '>' + children.map(({ type: t }) => t).join(',') : ''
+          }`,
+      ),
+    ).toEqual(b);
+  });
   test('metadata is parsed and becomes data', async () => {
     const mdast = u('root', [
       u('block', { meta: '{"key": "value"}' }, [
diff --git a/packages/myst-transforms/src/blocks.ts b/packages/myst-transforms/src/blocks.ts
index 1b5c6e3c7..03249bf8b 100644
--- a/packages/myst-transforms/src/blocks.ts
+++ b/packages/myst-transforms/src/blocks.ts
@@ -1,25 +1,32 @@
 import type { VFile } from 'vfile';
 import type { Plugin } from 'unified';
-import type { Node, Parent } from 'myst-spec';
-import { select, selectAll } from 'unist-util-select';
+import type { Node } from 'myst-spec';
+import { selectAll } from 'unist-util-select';
 import type { GenericNode, GenericParent } from 'myst-common';
 import { RuleId, fileError, normalizeLabel } from 'myst-common';
 import type { Code } from 'myst-spec-ext';
 
-export function blockNestingTransform(mdast: GenericParent) {
-  if (!select('block', mdast)) {
-    const blockNode = { type: 'block', children: mdast.children as Node[] };
-    (mdast as Parent).children = [blockNode];
+function nestBlock(tree: GenericParent): void {
+  const start = tree.children.findIndex((child) => (child.type as any) !== 'block');
+  if (start === -1) return;
+  const end = tree.children.findIndex((child, i) => (child.type as any) === 'block' && i > start);
+  if (end === -1) {
+    tree.children = [
+      ...tree.children.slice(0, start),
+      { type: 'block', children: tree.children.slice(start) as Node[] },
+    ];
     return;
   }
-  if ((mdast.children[0].type as any) !== 'block') {
-    // There are some blocks, but the first one is not!
-    // Lift the content until the first block into a block
-    const index = mdast.children.findIndex((child) => (child.type as any) === 'block');
-    const firstBlock = { type: 'block', children: mdast.children.slice(0, index) as Node[] };
-    const otherBlocks = mdast.children.slice(index) as Parent[];
-    (mdast as Parent).children = [firstBlock, ...otherBlocks];
-  }
+  tree.children = [
+    ...tree.children.slice(0, start),
+    { type: 'block', children: tree.children.slice(start, end) as Node[] },
+    ...tree.children.slice(end),
+  ];
+  nestBlock(tree);
+}
+
+export function blockNestingTransform(tree: GenericParent) {
+  nestBlock(tree);
 }
 
 export const blockNestingPlugin: Plugin<[], GenericParent, GenericParent> = () => (tree) => {