1 files changed, 62 insertions, 11 deletions

diff --git a/Map/Business/WmapParser/Entity.swift b/Map/Business/WmapParser/Entity.swift
index 40dd7ca..52f6e93 100644
--- a/Map/Business/WmapParser/Entity.swift
+++ b/Map/Business/WmapParser/Entity.swift
@@ -1,11 +1,51 @@
 struct Wmap {
-  protocol Entity {
-    var position: SourcePosition { get }
-    var range: SourceRange { get }
-    var syntaxElements: [SyntaxElement] { get }
+  enum Entity: Equatable {
+    case vertex(Vertex)
+    case edge(Edge)
+    case note(Note)
+    case stage(Stage)
+    case group(Group)
+    case inertia(Inertia)
+    case evolution(Evolution)
+
+    var position: SourcePosition {
+      switch self {
+      case .vertex(let vertex): return vertex.position
+      case .edge(let edge): return edge.position
+      case .note(let note): return note.position
+      case .stage(let stage): return stage.position
+      case .group(let group): return group.position
+      case .inertia(let inertia): return inertia.position
+      case .evolution(let evolution): return evolution.position
+      }
+    }
+
+    var range: SourceRange {
+      switch self {
+      case .vertex(let vertex): return vertex.range
+      case .edge(let edge): return edge.range
+      case .note(let note): return note.range
+      case .stage(let stage): return stage.range
+      case .group(let group): return group.range
+      case .inertia(let inertia): return inertia.range
+      case .evolution(let evolution): return evolution.range
+      }
+    }
+
+    var syntaxElements: [SyntaxElement] {
+      switch self {
+      case .vertex(let vertex): return vertex.syntaxElements
+      case .edge(let edge): return edge.syntaxElements
+      case .note(let note): return note.syntaxElements
+      case .stage(let stage): return stage.syntaxElements
+      case .group(let group): return group.syntaxElements
+      case .inertia(let inertia): return inertia.syntaxElements
+      case .evolution(let evolution): return evolution.syntaxElements
+      }
+    }
   }
 
-  struct Vertex: Entity {
+  struct Vertex: Equatable {
     let label: String
     let position: SourcePosition
     let coordinates: (Double, Double)
@@ -21,6 +61,12 @@ struct Wmap {
     let rightParenRange: SourceRange
     let shapeRange: SourceRange?
 
+    static func == (lhs: Vertex, rhs: Vertex) -> Bool {
+      return lhs.label == rhs.label && lhs.position == rhs.position
+        && lhs.coordinates.0 == rhs.coordinates.0 && lhs.coordinates.1 == rhs.coordinates.1
+        && lhs.shape == rhs.shape && lhs.range == rhs.range
+    }
+
     var syntaxElements: [SyntaxElement] {
       var elements: [SyntaxElement] = []
 
@@ -86,7 +132,7 @@ struct Wmap {
     }
   }
 
-  struct Edge: Entity {
+  struct Edge: Equatable {
     let from: String
     let to: String
     let isDirected: Bool
@@ -129,7 +175,7 @@ struct Wmap {
     }
   }
 
-  struct Note: Entity {
+  struct Note: Equatable {
     let coordinates: (Double, Double)
     let text: String
     let position: SourcePosition
@@ -144,6 +190,11 @@ struct Wmap {
     let rightParenRange: SourceRange
     let textRange: SourceRange
 
+    static func == (lhs: Note, rhs: Note) -> Bool {
+      return lhs.coordinates.0 == rhs.coordinates.0 && lhs.coordinates.1 == rhs.coordinates.1
+        && lhs.text == rhs.text && lhs.position == rhs.position && lhs.range == rhs.range
+    }
+
     var syntaxElements: [SyntaxElement] {
       var elements: [SyntaxElement] = []
 
@@ -207,7 +258,7 @@ struct Wmap {
     }
   }
 
-  struct Stage: Entity {
+  struct Stage: Equatable {
     let number: String  // i, ii, iii, iv
     let value: Double
     let position: SourcePosition
@@ -258,7 +309,7 @@ struct Wmap {
     }
   }
 
-  struct Group: Entity {
+  struct Group: Equatable {
     let vertices: [String]
     let position: SourcePosition
     let range: SourceRange
@@ -303,7 +354,7 @@ struct Wmap {
     }
   }
 
-  struct Inertia: Entity {
+  struct Inertia: Equatable {
     let vertex: String
     let position: SourcePosition
     let range: SourceRange
@@ -335,7 +386,7 @@ struct Wmap {
     }
   }
 
-  struct Evolution: Entity {
+  struct Evolution: Equatable {
     let vertex: String
     let isPositive: Bool
     let value: Double