feat: improve ir

packages/ir/src/index.ts

@@ -3,14 +3,16 @@ export interface Show {
 }
 
 export class GraphModel implements Show {
+  // TODO: Use sets instead of arrays
   readonly #nodes: GraphNode[] = []
   readonly #nodeMap: Map<string, GraphNode> = new Map()
   readonly #edges: GraphEdge[] = []
 
-  public constructor(
-    public readonly root: GraphNode,
-    public readonly idAttribute: string
-  ) {}
+  public readonly root: GraphNode
+
+  public constructor(public readonly idAttribute: string, rootTag: string) {
+    this.root = this.addNode(rootTag)
+  }
 
   public get nodes(): Readonly<GraphNode[]> {
     return this.#nodes
@@ -42,16 +44,52 @@ export class GraphModel implements Show {
     return this.getNearestIdentifiableNode(parent)
   }
 
-  public addNode(node: GraphNode) {
+  public addNode(tag: string) {
+    const node = new GraphNode(this, tag)
     this.#nodes.push(node)
     const id = this.getNodeId(node)
     if (id !== undefined) {
       this.#nodeMap.set(id, node)
     }
+    return node
+  }
+
+  public removeNode(node: GraphNode) {
+    if (node.model !== this) {
+      throw new Error('Node is not part of this model')
+    }
+    node.incomingEdges.forEach((edge) => this.removeEdge(edge))
+    node.outgoingEdges.forEach((edge) => this.removeEdge(edge))
+    const id = this.getNodeId(node)
+    if (id !== undefined) {
+      this.#nodeMap.delete(id)
+    }
+    node.children.forEach((child) => this.removeNode(child))
+    node.parent?.removeChild(node)
+    this.#nodes.splice(this.#nodes.indexOf(node), 1)
   }
 
-  public addEdge(edge: GraphEdge) {
+  public addEdge(tag: string, source: GraphNode, target: GraphNode) {
+    if (source.model !== this) {
+      throw new Error('Source node is not part of this model')
+    }
+    if (target.model !== this) {
+      throw new Error('Target node is not part of this model')
+    }
+    const edge = new GraphEdge(this, tag, source, target)
+    source.addOutgoingEdge(edge)
+    target.addOutgoingEdge(edge)
     this.#edges.push(edge)
+    return edge
+  }
+
+  public removeEdge(edge: GraphEdge) {
+    if (edge.model !== this) {
+      throw new Error('Edge is not part of this model')
+    }
+    edge.source.removeOutgoingEdge(edge)
+    edge.target.removeIncomingEdge(edge)
+    this.#edges.splice(this.#edges.indexOf(edge), 1)
   }
 
   public show(): string {
@@ -69,7 +107,10 @@ export class GraphNode implements Show {
   readonly #outgoingEdges = new Set<GraphEdge>()
   readonly #incomingEdges = new Set<GraphEdge>()
 
-  public constructor(public readonly tag: string) {}
+  public constructor(
+    public readonly model: GraphModel,
+    public readonly tag: string
+  ) {}
 
   public get parent(): GraphNode | undefined {
     return this.#parent
@@ -141,18 +182,30 @@ export class GraphNode implements Show {
   }
 
   public addIncomingEdge(edge: GraphEdge) {
+    if (edge.source !== this) {
+      throw new Error('Edge source is not this node')
+    }
     this.#incomingEdges.add(edge)
   }
 
   public removeIncomingEdge(edge: GraphEdge) {
+    if (edge.source !== this) {
+      throw new Error('Edge source is not this node')
+    }
     this.#incomingEdges.delete(edge)
   }
 
   public addOutgoingEdge(edge: GraphEdge) {
+    if (edge.target !== this) {
+      throw new Error('Edge target is not this node')
+    }
     this.#outgoingEdges.add(edge)
   }
 
   public removeOutgoingEdge(edge: GraphEdge) {
+    if (edge.target !== this) {
+      throw new Error('Edge target is not this node')
+    }
     this.#outgoingEdges.delete(edge)
   }
 
@@ -210,15 +263,14 @@ export interface NamespacedValue extends SimpleValue {
 }
 export type Value = SimpleValue | NamespacedValue
 
+// TODO: Add attributes to graph edge via delegate
 export class GraphEdge {
   public constructor(
+    public readonly model: GraphModel,
     public readonly tag: string,
     public readonly source: GraphNode,
     public readonly target: GraphNode
-  ) {
-    source.addOutgoingEdge(this)
-    target.addIncomingEdge(this)
-  }
+  ) {}
 }
 
 function createIndent(indent: number): string {

packages/uml-parser/src/index.ts

@@ -1,4 +1,3 @@
-import { GraphEdge } from '@cm2ml/ir'
 import type { GraphModel, GraphNode } from '@cm2ml/ir'
 import { compose, definePlugin } from '@cm2ml/plugin'
 import { XmiParser } from '@cm2ml/xmi-parser'
@@ -31,9 +30,9 @@ function refine(
   strict: boolean,
   greedyEdges: boolean
 ): GraphModel {
-  Stream.from(model.nodes)
-    .flatMap((node) => createEdges(node, model, strict, greedyEdges))
-    .forEach((edge) => model.addEdge(edge))
+  Stream.from(model.nodes).forEach((node) =>
+    createEdges(node, model, strict, greedyEdges)
+  )
   // console.log(
   //   model.edges
   //     .map((edge) => `${edge.source.tag} --${edge.tag}-> ${edge.target.tag}`)
@@ -49,80 +48,77 @@ function createEdges(
   model: GraphModel,
   strict: boolean,
   greedyEdges: boolean
-): Stream<GraphEdge> {
+) {
   if (greedyEdges) {
-    return createGreedyEdges(node, model)
+    createGreedyEdges(node, model)
+    return
   }
   if (ignoredTags.includes(node.tag)) {
-    return Stream.empty()
+    return
   }
   switch (node.tag) {
     case 'generalization':
-      return createGeneralizationEdges(node, model, strict)
+      createGeneralizationEdges(node, model, strict)
+      return
     case 'packagedElement':
-      return Stream.empty()
+      return
     default:
       if (strict) {
         throw new Error(`Unhandled tag: ${node.tag}`)
       }
-      return Stream.empty()
   }
 }
 
-function createGreedyEdges(
-  node: GraphNode,
-  model: GraphModel
-): Stream<GraphEdge> {
-  return Stream.from(node.attributes)
+function createGreedyEdges(node: GraphNode, model: GraphModel) {
+  Stream.from(node.attributes)
     .filter(([, { name }]) => name !== model.idAttribute)
-    .map(([_name, attribute]) => {
+    .forEach(([_name, attribute]) => {
       const source = model.getNearestIdentifiableNode(node)
       if (!source) {
-        return null
+        return
       }
       const attributeValue = attribute.value.literal
       const target = model.getNodeById(attributeValue)
       if (!target) {
-        return null
+        return
       }
       const tag = attribute.name === 'idref' ? node.tag : attribute.name
-      return new GraphEdge(tag, source, target)
+      model.addEdge(tag, source, target)
     })
-    .filterNonNull()
 }
 
 function createGeneralizationEdges(
   generalization: GraphNode,
   model: GraphModel,
   strict: boolean
-): Stream<GraphEdge> {
+) {
   const source = generalization.parent
   if (!source) {
     if (strict) {
       throw new Error('Expected parent of <generalization />')
     }
-    return Stream.empty()
+    return
   }
   const general = generalization.findChild((child) => child.tag === 'general')
   if (!general) {
     if (strict) {
       throw new Error('Expected <general /> child')
     }
-    return Stream.empty()
+    return
   }
   const targetId = general.getAttribute('idref')?.value.literal
   if (!targetId) {
     if (strict) {
       throw new Error('Expected idref attribute on <general />')
     }
-    return Stream.empty()
+    return
   }
   const target = model.getNodeById(targetId)
   if (!target) {
     if (strict) {
       throw new Error(`Expected node with id ${targetId}`)
     }
-    return Stream.empty()
+    return
   }
-  return Stream.fromSingle(new GraphEdge('generalization', source, target))
+  model.addEdge('generalization', source, target)
 }

packages/xmi-parser/src/index.ts

@@ -1,5 +1,5 @@
-import type { Attribute, Value } from '@cm2ml/ir'
-import { GraphModel, GraphNode } from '@cm2ml/ir'
+import type { Attribute, GraphNode, Value } from '@cm2ml/ir'
+import { GraphModel } from '@cm2ml/ir'
 import { definePlugin } from '@cm2ml/plugin'
 import { Stream } from '@yeger/streams'
 import { Element } from 'domhandler'
@@ -23,17 +23,10 @@ function parse(xmi: string, idAttribute: string): GraphModel {
   const document = parseDocument(xmi, {
     xmlMode: true,
   })
-  return getGraphModel(document, idAttribute)
+  return mapDocument(document, idAttribute)
 }
 
-function getGraphModel(document: Document, idAttribute: string): GraphModel {
-  const root = mapDocument(document)
-  const model = new GraphModel(root, idAttribute)
-  collectAllElements(root).forEach((node) => model.addNode(node))
-  return model
-}
-
-function mapDocument(document: Document) {
+function mapDocument(document: Document, idAttribute: string) {
   const elementChildren = Stream.from(document.childNodes)
     .map((node) => (isElement(node) ? node : null))
     .filterNonNull()
@@ -42,20 +35,29 @@ function mapDocument(document: Document) {
   if (elementChildren.length !== 1 || !root) {
     throw new Error('Expected exactly one root element')
   }
-  return mapElement(root)
+  const model = new GraphModel(idAttribute, root.tagName)
+  initNodeFromElement(model.root, root)
+  return model
 }
 
-function mapElement(element: Element): GraphNode {
-  const xmiElement = new GraphNode(element.tagName)
+function createNodeFromElement(model: GraphModel, element: Element): GraphNode {
+  const node = model.addNode(element.tagName)
+  initNodeFromElement(node, element)
+  return node
+}
+
+function initNodeFromElement(node: GraphNode, element: Element) {
   Stream.fromObject(element.attribs)
     .map(mapAttribute)
-    .forEach((attribute) => xmiElement.addAttribute(attribute, true))
+    .forEach((attribute) => node.addAttribute(attribute, true))
   Stream.from(element.childNodes)
-    .map((child) => (isElement(child) ? mapElement(child) : null))
+    .map((child) =>
+      isElement(child) ? createNodeFromElement(node.model, child) : null
+    )
     .filterNonNull()
-    .forEach((child) => xmiElement.addChild(child))
-  return xmiElement
+    .forEach((child) => node.addChild(child))
 }
+
 function mapAttribute([name, value]: [string, string]): Attribute {
   const xmiValue = mapValue(value)
   if (!name.includes(':')) {
@@ -76,9 +78,3 @@ function mapValue(value: string): Value {
 function isElement(node: Node): node is Element {
   return node.type === 'tag' || node instanceof Element
 }
-
-function collectAllElements(element: GraphNode): Stream<GraphNode> {
-  return Stream.from(element.children)
-    .flatMap(collectAllElements)
-    .append(element)
-}