feat: enhance Mermaid export with subgraphs, edge labels, node shapes and styling

.codegraphrc.json

@@ -0,0 +1,3 @@
+{
+  "embeddings": { "model": "bge-large" }
+}

src/cli.js

@@ -274,13 +274,15 @@ program
   .option('-T, --no-tests', 'Exclude test/spec files')
   .option('--include-tests', 'Include test/spec files (overrides excludeTests config)')
   .option('--min-confidence <score>', 'Minimum edge confidence threshold (default: 0.5)', '0.5')
+  .option('--direction <dir>', 'Flowchart direction for Mermaid: TB, LR, RL, BT', 'LR')
   .option('-o, --output <file>', 'Write to file instead of stdout')
   .action((opts) => {
     const db = openReadonlyOrFail(opts.db);
     const exportOpts = {
       fileLevel: !opts.functions,
       noTests: resolveNoTests(opts),
       minConfidence: parseFloat(opts.minConfidence),
+      direction: opts.direction,
     };
 
     let output;

src/export.js

@@ -125,19 +125,63 @@ export function exportDOT(db, opts = {}) {
   return lines.join('\n');
 }
 
+/** Escape double quotes for Mermaid labels. */
+function escapeLabel(label) {
+  return label.replace(/"/g, '#quot;');
+}
+
+/** Map node kind to Mermaid shape wrapper. */
+function mermaidShape(kind, label) {
+  const escaped = escapeLabel(label);
+  switch (kind) {
+    case 'function':
+    case 'method':
+      return `(["${escaped}"])`;
+    case 'class':
+    case 'interface':
+    case 'type':
+    case 'struct':
+    case 'enum':
+    case 'trait':
+    case 'record':
+      return `{{"${escaped}"}}`;
+    case 'module':
+      return `[["${escaped}"]]`;
+    default:
+      return `["${escaped}"]`;
+  }
+}
+
+/** Map node role to Mermaid style colors. */
+const ROLE_STYLES = {
+  entry: 'fill:#e8f5e9,stroke:#4caf50',
+  core: 'fill:#e3f2fd,stroke:#2196f3',
+  utility: 'fill:#f5f5f5,stroke:#9e9e9e',
+  dead: 'fill:#ffebee,stroke:#f44336',
+  leaf: 'fill:#fffde7,stroke:#fdd835',
+};
+
 /**
  * Export the dependency graph in Mermaid format.
  */
 export function exportMermaid(db, opts = {}) {
   const fileLevel = opts.fileLevel !== false;
   const noTests = opts.noTests || false;
   const minConf = opts.minConfidence ?? DEFAULT_MIN_CONFIDENCE;
-  const lines = ['graph LR'];
+  const direction = opts.direction || 'LR';
+  const lines = [`flowchart ${direction}`];
+
+  let nodeCounter = 0;
+  const nodeIdMap = new Map();
+  function nodeId(key) {
+    if (!nodeIdMap.has(key)) nodeIdMap.set(key, `n${nodeCounter++}`);
+    return nodeIdMap.get(key);
+  }
 
   if (fileLevel) {
     let edges = db
       .prepare(`
-      SELECT DISTINCT n1.file AS source, n2.file AS target
+      SELECT DISTINCT n1.file AS source, n2.file AS target, e.kind AS edge_kind
       FROM edges e
       JOIN nodes n1 ON e.source_id = n1.id
       JOIN nodes n2 ON e.target_id = n2.id
@@ -147,32 +191,133 @@ export function exportMermaid(db, opts = {}) {
       .all(minConf);
     if (noTests) edges = edges.filter((e) => !isTestFile(e.source) && !isTestFile(e.target));
 
+    // Collect all files referenced in edges
+    const allFiles = new Set();
     for (const { source, target } of edges) {
-      const s = source.replace(/[^a-zA-Z0-9]/g, '_');
-      const t = target.replace(/[^a-zA-Z0-9]/g, '_');
-      lines.push(`  ${s}["${source}"] --> ${t}["${target}"]`);
+      allFiles.add(source);
+      allFiles.add(target);
+    }
+
+    // Build directory groupings — try DB directory nodes first, fall back to path.dirname()
+    const dirs = new Map();
+    const hasDirectoryNodes =
+      db.prepare("SELECT COUNT(*) as c FROM nodes WHERE kind = 'directory'").get().c > 0;
+
+    if (hasDirectoryNodes) {
+      const dbDirs = db.prepare("SELECT id, name FROM nodes WHERE kind = 'directory'").all();
+      for (const d of dbDirs) {
+        const containedFiles = db
+          .prepare(`
+            SELECT n.name FROM edges e
+            JOIN nodes n ON e.target_id = n.id
+            WHERE e.source_id = ? AND e.kind = 'contains' AND n.kind = 'file'
+          `)
+          .all(d.id)
+          .map((r) => r.name)
+          .filter((f) => allFiles.has(f));
+        if (containedFiles.length > 0) dirs.set(d.name, containedFiles);
+      }
+    } else {
+      for (const file of allFiles) {
+        const dir = path.dirname(file) || '.';
+        if (!dirs.has(dir)) dirs.set(dir, []);
+        dirs.get(dir).push(file);
+      }
+    }
+
+    // Emit subgraphs
+    for (const [dir, files] of [...dirs].sort((a, b) => a[0].localeCompare(b[0]))) {
+      const sgId = dir.replace(/[^a-zA-Z0-9]/g, '_');
+      lines.push(`  subgraph ${sgId}["${escapeLabel(dir)}"]`);
+      for (const f of files) {
+        const nId = nodeId(f);
+        lines.push(`    ${nId}["${escapeLabel(path.basename(f))}"]`);
+      }
+      lines.push('  end');
+    }
+
+    // Deduplicate edges per source-target pair, collecting all distinct kinds
+    const edgeMap = new Map();
+    for (const { source, target, edge_kind } of edges) {
+      const key = `${source}|${target}`;
+      const label = edge_kind === 'imports-type' ? 'imports' : edge_kind;
+      if (!edgeMap.has(key)) edgeMap.set(key, { source, target, labels: new Set() });
+      edgeMap.get(key).labels.add(label);
+    }
+
+    for (const { source, target, labels } of edgeMap.values()) {
+      lines.push(`  ${nodeId(source)} -->|${[...labels].join(', ')}| ${nodeId(target)}`);
     }
   } else {
     let edges = db
       .prepare(`
-      SELECT n1.name AS source_name, n1.file AS source_file,
-             n2.name AS target_name, n2.file AS target_file
+      SELECT n1.name AS source_name, n1.kind AS source_kind, n1.file AS source_file,
+             n2.name AS target_name, n2.kind AS target_kind, n2.file AS target_file,
+             e.kind AS edge_kind
       FROM edges e
       JOIN nodes n1 ON e.source_id = n1.id
       JOIN nodes n2 ON e.target_id = n2.id
-      WHERE n1.kind IN ('function', 'method', 'class', 'interface', 'type', 'struct', 'enum', 'trait', 'record', 'module') AND n2.kind IN ('function', 'method', 'class', 'interface', 'type', 'struct', 'enum', 'trait', 'record', 'module')
-      AND e.kind = 'calls'
-      AND e.confidence >= ?
+      WHERE n1.kind IN ('function', 'method', 'class', 'interface', 'type', 'struct', 'enum', 'trait', 'record', 'module')
+        AND n2.kind IN ('function', 'method', 'class', 'interface', 'type', 'struct', 'enum', 'trait', 'record', 'module')
+        AND e.kind = 'calls'
+        AND e.confidence >= ?
     `)
       .all(minConf);
     if (noTests)
       edges = edges.filter((e) => !isTestFile(e.source_file) && !isTestFile(e.target_file));
 
+    // Group nodes by file for subgraphs
+    const fileNodes = new Map();
+    const nodeKinds = new Map();
+    for (const e of edges) {
+      const sKey = `${e.source_file}::${e.source_name}`;
+      const tKey = `${e.target_file}::${e.target_name}`;
+      nodeId(sKey);
+      nodeId(tKey);
+      nodeKinds.set(sKey, e.source_kind);
+      nodeKinds.set(tKey, e.target_kind);
+
+      if (!fileNodes.has(e.source_file)) fileNodes.set(e.source_file, new Map());
+      fileNodes.get(e.source_file).set(sKey, e.source_name);
+
+      if (!fileNodes.has(e.target_file)) fileNodes.set(e.target_file, new Map());
+      fileNodes.get(e.target_file).set(tKey, e.target_name);
+    }
+
+    // Emit subgraphs grouped by file
+    for (const [file, nodes] of [...fileNodes].sort((a, b) => a[0].localeCompare(b[0]))) {
+      const sgId = file.replace(/[^a-zA-Z0-9]/g, '_');
+      lines.push(`  subgraph ${sgId}["${escapeLabel(file)}"]`);
+      for (const [key, name] of nodes) {
+        const kind = nodeKinds.get(key);
+        lines.push(`    ${nodeId(key)}${mermaidShape(kind, name)}`);
+      }
+      lines.push('  end');
+    }
+
+    // Emit edges with labels
     for (const e of edges) {
-      const sId = `${e.source_file}_${e.source_name}`.replace(/[^a-zA-Z0-9]/g, '_');
-      const tId = `${e.target_file}_${e.target_name}`.replace(/[^a-zA-Z0-9]/g, '_');
-      lines.push(`  ${sId}["${e.source_name}"] --> ${tId}["${e.target_name}"]`);
+      const sId = nodeId(`${e.source_file}::${e.source_name}`);
+      const tId = nodeId(`${e.target_file}::${e.target_name}`);
+      lines.push(`  ${sId} -->|${e.edge_kind}| ${tId}`);
+    }
+
+    // Role styling — query roles for all referenced nodes
+    const allKeys = [...nodeIdMap.keys()];
+    const roleStyles = [];
+    for (const key of allKeys) {
+      const colonIdx = key.indexOf('::');
+      if (colonIdx === -1) continue;
+      const file = key.slice(0, colonIdx);
+      const name = key.slice(colonIdx + 2);
+      const row = db
+        .prepare('SELECT role FROM nodes WHERE file = ? AND name = ? AND role IS NOT NULL LIMIT 1')
+        .get(file, name);
+      if (row?.role && ROLE_STYLES[row.role]) {
+        roleStyles.push(`  style ${nodeIdMap.get(key)} ${ROLE_STYLES[row.role]}`);
+      }
     }
+    lines.push(...roleStyles);
   }
 
   return lines.join('\n');

tests/graph/export.test.js

@@ -43,17 +43,65 @@ describe('exportDOT', () => {
 });
 
 describe('exportMermaid', () => {
-  it('generates valid Mermaid syntax', () => {
+  it('generates valid Mermaid syntax with flowchart LR default', () => {
     const db = createTestDb();
     const a = insertNode(db, 'src/a.js', 'file', 'src/a.js', 0);
     const b = insertNode(db, 'src/b.js', 'file', 'src/b.js', 0);
     insertEdge(db, a, b, 'imports');
 
     const mermaid = exportMermaid(db);
-    expect(mermaid).toContain('graph LR');
+    expect(mermaid).toContain('flowchart LR');
     expect(mermaid).toContain('-->');
     db.close();
   });
+
+  it('uses custom direction option', () => {
+    const db = createTestDb();
+    const a = insertNode(db, 'src/a.js', 'file', 'src/a.js', 0);
+    const b = insertNode(db, 'src/b.js', 'file', 'src/b.js', 0);
+    insertEdge(db, a, b, 'imports');
+
+    const mermaid = exportMermaid(db, { direction: 'TB' });
+    expect(mermaid).toContain('flowchart TB');
+    db.close();
+  });
+
+  it('groups files into directory subgraphs', () => {
+    const db = createTestDb();
+    const a = insertNode(db, 'src/a.js', 'file', 'src/a.js', 0);
+    const b = insertNode(db, 'lib/b.js', 'file', 'lib/b.js', 0);
+    insertEdge(db, a, b, 'imports');
+
+    const mermaid = exportMermaid(db);
+    expect(mermaid).toContain('subgraph');
+    expect(mermaid).toContain('"src"');
+    expect(mermaid).toContain('"lib"');
+    expect(mermaid).toContain('end');
+    db.close();
+  });
+
+  it('adds edge labels from edge kind', () => {
+    const db = createTestDb();
+    const a = insertNode(db, 'src/a.js', 'file', 'src/a.js', 0);
+    const b = insertNode(db, 'src/b.js', 'file', 'src/b.js', 0);
+    insertEdge(db, a, b, 'imports');
+
+    const mermaid = exportMermaid(db);
+    expect(mermaid).toContain('-->|imports|');
+    db.close();
+  });
+
+  it('collapses imports-type to imports label', () => {
+    const db = createTestDb();
+    const a = insertNode(db, 'src/a.js', 'file', 'src/a.js', 0);
+    const b = insertNode(db, 'src/b.js', 'file', 'src/b.js', 0);
+    insertEdge(db, a, b, 'imports-type');
+
+    const mermaid = exportMermaid(db);
+    expect(mermaid).toContain('-->|imports|');
+    expect(mermaid).not.toContain('imports-type');
+    db.close();
+  });
 });
 
 describe('exportDOT — function-level', () => {
@@ -107,12 +155,86 @@ describe('exportMermaid — function-level', () => {
     insertEdge(db, fnA, fnB, 'calls');
 
     const mermaid = exportMermaid(db, { fileLevel: false });
-    expect(mermaid).toContain('graph LR');
+    expect(mermaid).toContain('flowchart LR');
     expect(mermaid).toContain('doWork');
     expect(mermaid).toContain('helper');
     expect(mermaid).toContain('-->');
     db.close();
   });
+
+  it('uses stadium shape for functions', () => {
+    const db = createTestDb();
+    const fnA = insertNode(db, 'doWork', 'function', 'src/a.js', 5);
+    const fnB = insertNode(db, 'helper', 'function', 'src/b.js', 10);
+    insertEdge(db, fnA, fnB, 'calls');
+
+    const mermaid = exportMermaid(db, { fileLevel: false });
+    expect(mermaid).toContain('(["doWork"])');
+    expect(mermaid).toContain('(["helper"])');
+    db.close();
+  });
+
+  it('uses hexagon shape for classes', () => {
+    const db = createTestDb();
+    const cls = insertNode(db, 'MyClass', 'class', 'src/a.js', 5);
+    const fn = insertNode(db, 'helper', 'function', 'src/b.js', 10);
+    insertEdge(db, cls, fn, 'calls');
+
+    const mermaid = exportMermaid(db, { fileLevel: false });
+    expect(mermaid).toContain('{{"MyClass"}}');
+    db.close();
+  });
+
+  it('uses subroutine shape for modules', () => {
+    const db = createTestDb();
+    const mod = insertNode(db, 'MyModule', 'module', 'src/a.js', 5);
+    const fn = insertNode(db, 'helper', 'function', 'src/b.js', 10);
+    insertEdge(db, mod, fn, 'calls');
+
+    const mermaid = exportMermaid(db, { fileLevel: false });
+    expect(mermaid).toContain('[["MyModule"]]');
+    db.close();
+  });
+
+  it('adds edge labels for calls', () => {
+    const db = createTestDb();
+    const fnA = insertNode(db, 'doWork', 'function', 'src/a.js', 5);
+    const fnB = insertNode(db, 'helper', 'function', 'src/b.js', 10);
+    insertEdge(db, fnA, fnB, 'calls');
+
+    const mermaid = exportMermaid(db, { fileLevel: false });
+    expect(mermaid).toContain('-->|calls|');
+    db.close();
+  });
+
+  it('groups functions by file into subgraphs', () => {
+    const db = createTestDb();
+    const fnA = insertNode(db, 'doWork', 'function', 'src/a.js', 5);
+    const fnB = insertNode(db, 'helper', 'function', 'src/b.js', 10);
+    insertEdge(db, fnA, fnB, 'calls');
+
+    const mermaid = exportMermaid(db, { fileLevel: false });
+    expect(mermaid).toContain('subgraph');
+    expect(mermaid).toContain('"src/a.js"');
+    expect(mermaid).toContain('"src/b.js"');
+    expect(mermaid).toContain('end');
+    db.close();
+  });
+
+  it('applies role styling', () => {
+    const db = createTestDb();
+    const fnA = insertNode(db, 'doWork', 'function', 'src/a.js', 5);
+    const fnB = insertNode(db, 'helper', 'function', 'src/b.js', 10);
+    // Add role to the nodes
+    db.prepare('UPDATE nodes SET role = ? WHERE id = ?').run('entry', fnA);
+    db.prepare('UPDATE nodes SET role = ? WHERE id = ?').run('utility', fnB);
+    insertEdge(db, fnA, fnB, 'calls');
+
+    const mermaid = exportMermaid(db, { fileLevel: false });
+    expect(mermaid).toContain('fill:#e8f5e9,stroke:#4caf50');
+    expect(mermaid).toContain('fill:#f5f5f5,stroke:#9e9e9e');
+    db.close();
+  });
 });
 
 describe('exportJSON', () => {

tests/integration/cli.test.js

@@ -130,9 +130,9 @@ describe('CLI smoke tests', () => {
   });
 
   // ─── Export (Mermaid) ────────────────────────────────────────────────
-  test('export -f mermaid outputs graph LR', () => {
+  test('export -f mermaid outputs flowchart LR', () => {
     const out = run('export', '--db', dbPath, '-f', 'mermaid');
-    expect(out).toContain('graph LR');
+    expect(out).toContain('flowchart LR');
   });
 
   // ─── Export (JSON) ───────────────────────────────────────────────────