feat: draw edges

Signed-off-by: Andres Correa Casablanca <castarco@coderspirit.xyz>

Author: castarco

src/BeautifulTree.tsx

@@ -1,6 +1,6 @@
+import { edgesIterator, postOrderIterator } from './traversal'
 import type { Tree } from './types'
 import { computeLeftShiftLayout } from './layouts'
-import { postOrderIterator } from './traversal'
 
 export interface BeautifulTreeProps {
 	readonly id: string
@@ -18,10 +18,11 @@ export function BeautifulTree({
 	tree,
 }: Readonly<BeautifulTreeProps>): JSX.Element {
 	const { tree: treeWithLayout, maxX, maxY } = computeLeftShiftLayout(tree)
-	const orderedNodes = [...postOrderIterator(treeWithLayout)]
-
 	const { width, height, sizeUnit = 'px' } = svgProps
 
+	const xDivisor = maxX + 2
+	const yDivisor = maxY + 2
+
 	return (
 		<svg
 			xmlns="http://www.w3.org/2000/svg"
@@ -33,15 +34,29 @@ export function BeautifulTree({
 			}}
 			className={'beautiful-tree-react'}
 		>
-			{orderedNodes.map((node, idx) => {
+			{/* TODO: introduce edge "styles" (straight, cornered, curved..., plus CSS styles) */}
+			{[...edgesIterator(treeWithLayout)].map((edge, idx) => {
+				return (
+					<line
+						key={`${id}-edge-${idx}`}
+						x1={((edge.start.x + 1) * width) / xDivisor}
+						y1={((edge.start.y + 1) * height) / yDivisor}
+						x2={((edge.end.x + 1) * width) / xDivisor}
+						y2={((edge.end.y + 1) * height) / yDivisor}
+						stroke="black"
+					/>
+				)
+			})}
+
+			{[...postOrderIterator(treeWithLayout)].map((node, idx) => {
 				const aX = node.meta.abstractPosition.x
 				const aY = node.meta.abstractPosition.y
 				return (
 					<circle
 						key={`${id}-node-${idx}`}
 						className={'beautiful-tree-node'}
-						cx={((aX + 1) * width) / (maxX + 2)}
-						cy={((aY + 1) * height) / (maxY + 2)}
+						cx={((aX + 1) * width) / xDivisor}
+						cy={((aY + 1) * height) / yDivisor}
 						stroke="black"
 						fill="white"
 						r="5"

src/tests/traversal.test.ts

@@ -1,5 +1,6 @@
 import { describe, expect, it } from 'vitest'
-import { postOrderIterator } from '../traversal'
+import { edgesIterator, postOrderIterator } from '../traversal'
+import type { TreeWithLayout } from '../types'
 
 describe('postOrderIterator', () => {
 	const testTree = {
@@ -38,8 +39,6 @@ describe('postOrderIterator', () => {
 
 	it('traverses tree in post-order', () => {
 		const result = postOrderIterator(testTree)
-
-		console.log('result:')
 		const expandedResult = [...result]
 
 		expect(expandedResult).toEqual([
@@ -52,3 +51,88 @@ describe('postOrderIterator', () => {
 		])
 	})
 })
+
+describe('edgesIterator', () => {
+	const testTree: TreeWithLayout = {
+		data: { v: 42 },
+		children: [
+			{
+				edgeData: { e: 100 },
+				node: {
+					data: { v: 43 },
+					children: [
+						{
+							edgeData: { e: 102 },
+							node: {
+								data: { v: 45 },
+								meta: {
+									isRoot: false,
+									isLeaf: true,
+									abstractPosition: { x: 0, y: 2 },
+								},
+							},
+						},
+					],
+					meta: {
+						isRoot: false,
+						isLeaf: false,
+						abstractPosition: { x: 0, y: 1 },
+					},
+				},
+			},
+			{
+				edgeData: { e: 101 },
+				node: {
+					data: { v: 44 },
+					children: [
+						{
+							edgeData: { e: 103 },
+							node: {
+								data: { v: 46 },
+								meta: {
+									isRoot: false,
+									isLeaf: true,
+									abstractPosition: { x: 1, y: 2 },
+								},
+							},
+						},
+						{
+							edgeData: { e: 104 },
+							node: {
+								data: { v: 47 },
+								meta: {
+									isRoot: false,
+									isLeaf: true,
+									abstractPosition: { x: 2, y: 2 },
+								},
+							},
+						},
+					],
+					meta: {
+						isRoot: false,
+						isLeaf: false,
+						abstractPosition: { x: 1, y: 1 },
+					},
+				},
+			},
+		],
+		meta: {
+			isRoot: true,
+			isLeaf: false,
+			abstractPosition: { x: 0, y: 0 },
+		},
+	}
+
+	it('traverses edges in pre-order', () => {
+		const result = edgesIterator(testTree)
+		const expandedResult = [...result]
+
+		expect(expandedResult).toEqual([
+			{ start: { x: 0, y: 0 }, end: { x: 0, y: 1 }, edgeData: { e: 100 } },
+			{ start: { x: 0, y: 1 }, end: { x: 0, y: 2 }, edgeData: { e: 102 } },
+			{ start: { x: 0, y: 0 }, end: { x: 1, y: 1 }, edgeData: { e: 101 } },
+			{ start: { x: 1, y: 1 }, end: { x: 1, y: 2 }, edgeData: { e: 103 } },
+			{ start: { x: 1, y: 1 }, end: { x: 2, y: 2 }, edgeData: { e: 104 } },
+		])
+	})
+})

src/traversal.ts

@@ -1,16 +1,29 @@
-import type { Tree, TreeWithLayout } from './types'
+import type { Edge, Tree, TreeWithLayout } from './types'
 
 export function* postOrderIterator<T extends Tree | TreeWithLayout = Tree>(
-	v: Readonly<T>,
+	tree: Readonly<T>,
 ): Generator<Omit<Readonly<T>, 'children'>, void> {
-	for (const child of v.children ?? []) {
+	for (const child of tree.children ?? []) {
 		yield* postOrderIterator<T>(child.node as T)
 	}
 
 	// We'll discard children, but we don't directly access v.data because we
 	// might also want to access v.layout in case it exists.
 
 	// eslint-disable-next-line @typescript-eslint/no-unused-vars
-	const { children, ...vv } = v
+	const { children, ...vv } = tree
 	yield vv
 }
+
+export function* edgesIterator(
+	tree: Readonly<TreeWithLayout>,
+): Generator<Readonly<Edge>, void> {
+	for (const child of tree.children ?? []) {
+		yield {
+			start: tree.meta.abstractPosition,
+			end: child.node.meta.abstractPosition,
+			edgeData: child.edgeData,
+		}
+		yield* edgesIterator(child.node)
+	}
+}

src/types.ts

@@ -12,7 +12,9 @@ export type Tree = Node & {
 }
 
 export type TreeWithLayout = Node & {
-	readonly children?: Required<TreeChild<TreeWithLayout>>[] | undefined
+	readonly children?:
+		| Readonly<Required<TreeChild<TreeWithLayout>>[]>
+		| undefined
 	readonly meta: {
 		readonly isRoot: boolean
 		readonly isLeaf: boolean
@@ -22,3 +24,15 @@ export type TreeWithLayout = Node & {
 		}
 	}
 }
+
+export interface Edge {
+	readonly start: {
+		readonly x: number
+		readonly y: number
+	}
+	readonly end: {
+		readonly x: number
+		readonly y: number
+	}
+	readonly edgeData?: Readonly<Record<string, unknown>> | undefined
+}