From cd64de40267277b8cf82ca5c6580e66852cc1e13 Mon Sep 17 00:00:00 2001
From: gmega <giuliano.mega@gmail.com>
Date: Fri, 25 Aug 2023 15:46:38 -0300
Subject: [PATCH] factor out message type macro into more general type id macro

---
 swarmsim/codex/dhttracker.nim | 28 ++++++------
 swarmsim/engine/message.nim   | 51 +--------------------
 swarmsim/engine/peer.nim      |  6 +--
 swarmsim/engine/protocol.nim  | 14 +++---
 swarmsim/engine/types.nim     |  9 ++--
 swarmsim/lib/withtypeid.nim   | 86 +++++++++++++++++++++++++++++++++++
 tests/all_tests.nim           |  2 +-
 tests/codex/tdhttracker.nim   |  2 +-
 tests/engine/tmessage.nim     | 20 --------
 tests/engine/tnetwork.nim     |  5 +-
 tests/engine/tpeer.nim        | 13 +++---
 tests/helpers/inbox.nim       |  3 ++
 tests/helpers/testpeer.nim    | 26 +++++++++++
 tests/helpers/types.nim       |  9 ++++
 tests/lib/twithtypeid.nim     | 38 ++++++++++++++++
 15 files changed, 203 insertions(+), 109 deletions(-)
 create mode 100644 swarmsim/lib/withtypeid.nim
 delete mode 100644 tests/engine/tmessage.nim
 create mode 100644 tests/helpers/testpeer.nim
 create mode 100644 tests/helpers/types.nim
 create mode 100644 tests/lib/twithtypeid.nim

diff --git a/swarmsim/codex/dhttracker.nim b/swarmsim/codex/dhttracker.nim
index 4681a52..8d57cc6 100644
--- a/swarmsim/codex/dhttracker.nim
+++ b/swarmsim/codex/dhttracker.nim
@@ -18,26 +18,25 @@ type
 
 type ArrayShuffler = proc (arr: var seq[PeerDescriptor]): void
 
-type
-  DHTTracker* = ref object of Protocol
-    peerExpiration*: Duration
-    maxPeers*: uint
-    peers: OrderedTable[int, PeerDescriptor]
-    shuffler: ArrayShuffler
-
-
-  ExpirationTimer* = ref object of SchedulableEvent
-    peerId*: int
-    tracker: DHTTracker
-
-typedMessage:
+withTypeId:
   type
+    DHTTracker* = ref object of Protocol
+      peerExpiration*: Duration
+      maxPeers*: uint
+      peers: OrderedTable[int, PeerDescriptor]
+      shuffler: ArrayShuffler
+
     PeerAnnouncement* = ref object of Message
       peerId*: int
 
     SampleSwarm* = ref object of Message
       numPeers: uint
 
+type
+  ExpirationTimer* = ref object of SchedulableEvent
+    peerId*: int
+    tracker: DHTTracker
+
 let RandomShuffler = proc (arr: var seq[PeerDescriptor]) =
   discard arr.nextPermutation()
 
@@ -51,12 +50,11 @@ proc new*(
 ): DHTTracker =
   DHTTracker(
     # This should in general be safe as those are always positive.
-    id: "DHTTracker",
     peerExpiration: peerExpiration,
     maxPeers: maxPeers,
     shuffler: shuffler,
     peers: initOrderedTable[int, PeerDescriptor](),
-    messageTypes: @[PeerAnnouncement.messageType, SampleSwarm.messageType]
+    messageTypes: @[PeerAnnouncement.typeId, SampleSwarm.typeId]
   )
 
 proc peers*(self: DHTTracker): seq[PeerDescriptor] = self.peers.values.toSeq()
diff --git a/swarmsim/engine/message.nim b/swarmsim/engine/message.nim
index 3afd341..1aac80d 100644
--- a/swarmsim/engine/message.nim
+++ b/swarmsim/engine/message.nim
@@ -1,54 +1,7 @@
-import options
-import macros
-
 import ./types
 
-method `messageType`*(self: Message): string {.base.} =
-  raise newException(CatchableError, "Method without implementation override")
-
-method `messageType`*(self: FreelyTypedMessage): string = self.messageType
+method `messageType`*(self: Message): string {.base.} = self.typeId
 
 proc allMessages*(self: type Message): string = "*"
 
-func typeName(typeDef: NimNode): Option[NimNode] =
-  expectKind typeDef, nnkTypeDef
-
-  return if typeDef[0].kind == nnkIdent:
-    typeDef[0].some
-  elif typeDef[0].kind == nnkPostfix:
-    typeDef[0][1].some
-  else:
-    none(NimNode)
-
-macro typedMessage*(body: untyped): untyped =
-  expectKind body, nnkStmtList
-  expectKind body[0], nnkTypeSection
-
-  for statement in body[0]:
-    if statement.kind != nnkTypeDef:
-      continue
-
-    let maybeTypename = typeName(statement)
-    if maybeTypename.isNone:
-      error("unable to get type name from AST. Sorry.")
-
-    let typeIdent = maybeTypename.get
-    let typeName = newLit(typeIdent.strVal)
-
-    let typeProc = quote do:
-      proc messageType*(self: type `typeIdent`): string = `typeName`
-
-    let instanceProc = quote do:
-      method messageType*(self: `typeIdent`): string = `typeIdent`.messageType
-
-    # We replace the proc name with a quoted symbol so it turns into a
-    # getter.
-    typeProc[0][1] = newTree(nnkAccQuoted, typeProc[0][1])
-    instanceProc[0][1] = newTree(nnkAccQuoted, instanceProc[0][1])
-
-    body.add(typeProc)
-    body.add(instanceProc)
-
-  return body
-
-export Message, FreelyTypedMessage
+export Message
diff --git a/swarmsim/engine/peer.nim b/swarmsim/engine/peer.nim
index 3fbf45d..618263d 100644
--- a/swarmsim/engine/peer.nim
+++ b/swarmsim/engine/peer.nim
@@ -21,8 +21,8 @@ proc getProtocol*(self: Peer, id: string): Option[Protocol] =
 
   none(Protocol)
 
-proc addProtocol*(self: Peer, protocol: Protocol): void =
-  self.protocols[protocol.id] = protocol
+proc addProtocol*[T: Protocol](self: Peer, protocol: T): void =
+  self.protocols[protocol.protocolId] = protocol
 
 proc deliverForType(self: Peer, messageType: string, message: Message,
     engine: EventDrivenEngine, network: Network): void =
@@ -44,7 +44,7 @@ proc initPeer*(self: Peer, protocols: seq[Protocol],
   for protocol in protocols:
     let protocol = protocol # https://github.com/nim-lang/Nim/issues/16740
 
-    self.protocols[protocol.id] = protocol
+    self.protocols[protocol.protocolId] = protocol
     protocol.messageTypes.apply(proc (m: string): void =
       self.dispatch.add(m, protocol))
 
diff --git a/swarmsim/engine/protocol.nim b/swarmsim/engine/protocol.nim
index 1420465..7d01dd9 100644
--- a/swarmsim/engine/protocol.nim
+++ b/swarmsim/engine/protocol.nim
@@ -1,5 +1,3 @@
-import typetraits
-
 import ./types
 import ./eventdrivenengine
 
@@ -7,8 +5,12 @@ export eventdrivenengine
 export Protocol
 export Message
 
-method deliver*(self: Protocol, message: Message, engine: EventDrivenEngine,
-    network: Network): void {.base.} =
-  raise newException(CatchableError, "Method without implementation override")
+method `protocolId`*(self: Protocol): string {.base.} = self.typeId
 
-proc protocolName*[T: Protocol](self: type T): string = name(T)
+method deliver*(
+  self: Protocol,
+  message: Message,
+  engine: EventDrivenEngine,
+  network: Network
+): void {.base.} =
+  raise newException(CatchableError, "Method without implementation override")
diff --git a/swarmsim/engine/types.nim b/swarmsim/engine/types.nim
index e742f63..9dbfcc4 100644
--- a/swarmsim/engine/types.nim
+++ b/swarmsim/engine/types.nim
@@ -4,11 +4,13 @@ import std/sets
 import std/options
 import std/random
 
+import ../lib/withtypeid
 import ../lib/multitable
 
 export heapqueue
 export option
 export random
+export withtypeid
 
 type
   SchedulableEvent* = ref object of RootObj
@@ -32,7 +34,6 @@ type
     ## A `Protocol` defines a P2P protocol. It handles messages meant for it,
     ## keeps internal state, and may expose services to other `Protocol`s within
     ## the same `Peer`.
-    id*: string
     messageTypes*: seq[string]
 
 type
@@ -48,11 +49,6 @@ type
     sender*: Option[Peer] = none(Peer)
     receiver*: Peer
 
-  FreelyTypedMessage* = ref object of Message
-    ## A `FreelyTypedMessage` is a `Message` that can be of any type.
-    ##
-    messageType*: string
-
 type
   Network* = ref object of RootObj
     ## A `Network` is a collection of `Peer`s that can communicate with each
@@ -61,3 +57,4 @@ type
     engine*: EventDrivenEngine
     defaultLinkDelay*: uint64
     peers*: HashSet[Peer] # TODO: use an array
+
diff --git a/swarmsim/lib/withtypeid.nim b/swarmsim/lib/withtypeid.nim
new file mode 100644
index 0000000..22f4dde
--- /dev/null
+++ b/swarmsim/lib/withtypeid.nim
@@ -0,0 +1,86 @@
+
+## This package adds a very basic interface and a macro to annotate and query
+## types about their ids at runtime. Type information is queried over a method
+## and uses dynamic dispatch, which means you can always recover the actual
+## type of an object, even as it is upcasted to more general types.
+##
+## This is a stopgap measure to allow us to, for instance, register dispatchers
+## based on type information, and do type equality comparisons.
+##
+## NB. This is very naively implemented right now, and won't ensure by a long
+##   shot that type IDs - which are currently just the type's name - are unique.
+##   If this proves to be a worthwhile effort, however, it would be possible to
+##   extend this to use hashes (e.g. signatureHash) or a global counter, and have
+##   a separate "typeName" attribute to query the actual name of the type.
+##
+
+import options
+import macros
+
+method `typeId`*(self: RootObj): string {.base.} =
+  ## Returns the type id of an object. This is currently a string and
+  ## conflates with the type's human-readable name, but the only hard
+  ## requirement is that this is a hashable object and has well-defined
+  ## identity semantics (==).
+  ##
+  ## If a type is not created with the `withTypeId` macro, then the method will
+  ## raise an exception unless manually overriden by subtypes.
+  raise (ref Defect)(msg: "Type has not been annotated with `withTypeId`.")
+
+method `typeId`*(self: ref RootObj): string {.base.} =
+  raise (ref Defect)(msg: "Type has not been annotated with `withTypeId`.")
+
+func typeName(typeDef: NimNode): Option[NimNode] =
+  expectKind typeDef, nnkTypeDef
+
+  return if typeDef[0].kind == nnkIdent:
+    typeDef[0].some
+  elif typeDef[0].kind == nnkPostfix:
+    typeDef[0][1].some
+  else:
+    none(NimNode)
+
+macro withTypeId*(body: untyped): untyped =
+  ## Creates a type with a `typeId` method and a proc bound to the type
+  ## itself which return the type's name.
+  runnableExamples:
+    withTypeId:
+      type
+        Foo = object of RootObj
+        Bar* = object of RootObj
+
+    doAssert Foo.typeId == "Foo"
+    doAssert Bar.typeId == "Bar"
+
+    doAssert Foo().typeId == "Foo"
+    doAssert Bar().typeId == "Bar"
+
+  expectKind body, nnkStmtList
+  expectKind body[0], nnkTypeSection
+
+  for statement in body[0]:
+    if statement.kind != nnkTypeDef:
+      continue
+
+    let maybeTypename = typeName(statement)
+    if maybeTypename.isNone:
+      error("unable to get type name from AST. Sorry.")
+
+    let typeIdent = maybeTypename.get
+    let typeName = newLit(typeIdent.strVal)
+
+    let typeProc = quote do:
+      proc typeId*(self: type `typeIdent`): string = `typeName`
+
+    let instanceProc = quote do:
+      method typeId*(self: `typeIdent`): string = `typeIdent`.typeId
+
+    # We replace the proc name with a quoted symbol so it turns into a
+    # getter.
+    typeProc[0][1] = newTree(nnkAccQuoted, typeProc[0][1])
+    instanceProc[0][1] = newTree(nnkAccQuoted, instanceProc[0][1])
+
+    body.add(typeProc)
+    body.add(instanceProc)
+
+  return body
diff --git a/tests/all_tests.nim b/tests/all_tests.nim
index 0a8dc66..2a31e9f 100644
--- a/tests/all_tests.nim
+++ b/tests/all_tests.nim
@@ -2,9 +2,9 @@ import engine/teventdrivenengine
 import engine/tschedulableevent
 import engine/tnetwork
 import engine/tpeer
-import engine/tmessage
 import codex/tdhttracker
 import lib/tmultitable
+import lib/twithtypeid
 
 {.warning[UnusedImport]: off.}
 
diff --git a/tests/codex/tdhttracker.nim b/tests/codex/tdhttracker.nim
index f42128f..bc798d0 100644
--- a/tests/codex/tdhttracker.nim
+++ b/tests/codex/tdhttracker.nim
@@ -13,7 +13,7 @@ import pkg/swarmsim/timeutils
 
 proc getPeerArray(tracker: Peer): seq[PeerDescriptor] =
   DHTTracker(
-      tracker.getProtocol(DHTTracker.protocolName).get()).peers
+      tracker.getProtocol(DHTTracker.typeId).get()).peers
 
 proc getPeerIdArray(tracker: Peer): seq[int] =
   getPeerArray(tracker).map(p => p.peerId)
diff --git a/tests/engine/tmessage.nim b/tests/engine/tmessage.nim
deleted file mode 100644
index 3f7d80f..0000000
--- a/tests/engine/tmessage.nim
+++ /dev/null
@@ -1,20 +0,0 @@
-import unittest
-
-import swarmsim/engine/message
-
-typedMessage:
-  type
-    PeerAnnouncement* = object of Message
-      peerId*: int
-
-    PrivateMessage = object of Message
-
-suite "message":
-  test "should automatically generate a type string for typedMessage types":
-    check(PeerAnnouncement.messageType == "PeerAnnouncement")
-    check(PrivateMessage.messageType == "PrivateMessage")
-
-  test "should automatically generate a type string for typedMessage instances":
-    check(PeerAnnouncement(peerId: 1).messageType == "PeerAnnouncement")
-    check(PrivateMessage().messageType == "PrivateMessage")
-
diff --git a/tests/engine/tnetwork.nim b/tests/engine/tnetwork.nim
index 0534d7a..74f2fe6 100644
--- a/tests/engine/tnetwork.nim
+++ b/tests/engine/tnetwork.nim
@@ -7,14 +7,15 @@ import swarmsim/engine/peer
 import swarmsim/engine/protocol
 
 import ../helpers/inbox
+import ../helpers/types
 
 suite "network":
   test "should dispatch message to the correct peer":
 
     let engine = EventDrivenEngine()
 
-    let i1 = Inbox(id: "inbox", messageTypes: @["m"])
-    let i2 = Inbox(id: "inbox", messageTypes: @["m"])
+    let i1 = Inbox(protocolId: "inbox", messageTypes: @["m"])
+    let i2 = Inbox(protocolId: "inbox", messageTypes: @["m"])
 
     let p1 = Peer.new(protocols = @[Protocol i1])
     let p2 = Peer.new(protocols = @[Protocol i2])
diff --git a/tests/engine/tpeer.nim b/tests/engine/tpeer.nim
index 447c882..64d0b02 100644
--- a/tests/engine/tpeer.nim
+++ b/tests/engine/tpeer.nim
@@ -7,6 +7,7 @@ import swarmsim/engine/peer
 import swarmsim/engine/message
 
 import ../helpers/inbox
+import ../helpers/types
 
 # We need this here as otherwise for some reason the nim compiler trips.
 proc `$`*(m: Message): string = repr m
@@ -32,8 +33,8 @@ suite "peer":
     check(not peerSet.contains(p1))
 
   test "should dispatch message to correct protocol":
-    let i1 = Inbox(id: "protocol1", messageTypes: @["m1"])
-    let i2 = Inbox(id: "protocol2", messageTypes: @["m2"])
+    let i1 = Inbox(protocolId: "protocol1", messageTypes: @["m1"])
+    let i2 = Inbox(protocolId: "protocol2", messageTypes: @["m2"])
 
     let peer = Peer.new(protocols = @[Protocol i1, i2])
 
@@ -51,8 +52,8 @@ suite "peer":
     check(i2.messages == @[m2])
 
   test "should dispatch a message to multiple protocols if they are listening on the same message type":
-    let i1 = Inbox(id: "protocol1", messageTypes: @["m1"])
-    let i2 = Inbox(id: "protocol2", messageTypes: @["m1"])
+    let i1 = Inbox(protocolId: "protocol1", messageTypes: @["m1"])
+    let i2 = Inbox(protocolId: "protocol2", messageTypes: @["m1"])
 
     let peer = Peer.new(protocols = @[Protocol i1, i2])
 
@@ -64,7 +65,7 @@ suite "peer":
     check(i2.messages == @[m1])
 
   test "should allow protocol to listen on multiple message types":
-    let i1 = Inbox(id: "protocol1", messageTypes: @["m1", "m2"])
+    let i1 = Inbox(protocolId: "protocol1", messageTypes: @["m1", "m2"])
 
     let peer = Peer.new(protocols = @[Protocol i1])
 
@@ -80,7 +81,7 @@ suite "peer":
 
 
   test "should deliver all message types when listening to Message.allMessages":
-    let i1 = Inbox(id: "protocol1", messageTypes: @[Message.allMessages])
+    let i1 = Inbox(protocolId: "protocol1", messageTypes: @[Message.allMessages])
 
     let peer = Peer.new(protocols = @[Protocol i1])
 
diff --git a/tests/helpers/inbox.nim b/tests/helpers/inbox.nim
index ea72197..8ad9527 100644
--- a/tests/helpers/inbox.nim
+++ b/tests/helpers/inbox.nim
@@ -5,6 +5,7 @@ import swarmsim/engine/network
 
 type
   Inbox* = ref object of Protocol
+    protocolId*: string
     messages*: seq[Message]
 
 method deliver*(
@@ -15,6 +16,8 @@ method deliver*(
 ) =
   self.messages.add(message)
 
+method `protocolId`*(self: Inbox): string = self.protocolId
+
 export Message
 export peer
 export protocol
diff --git a/tests/helpers/testpeer.nim b/tests/helpers/testpeer.nim
new file mode 100644
index 0000000..b284bcd
--- /dev/null
+++ b/tests/helpers/testpeer.nim
@@ -0,0 +1,26 @@
+import std/options
+import std/random
+
+import swarmsim/engine
+import swarmsim/engine/peer
+
+import ./inbox
+
+type TestPeer* = ref object of Peer
+  network: Network
+
+proc new*(
+  t: typedesc[TestPeer],
+  network: Network,
+  peerId: Option[int] = none(int),
+): TestPeer =
+  let peer: TestPeer = TestPeer(network: network)
+  discard peer.initPeer(protocols = @[Protocol Inbox()])
+  peer
+
+proc inbox*(peer: TestPeer): Inbox =
+  Inbox peer.getProtocol(Inbox.protocolName).get()
+
+proc send*(self: TestPeer, msg: Message): ScheduledEvent =
+  msg.sender = Peer(self).some
+  self.network.send(msg)
diff --git a/tests/helpers/types.nim b/tests/helpers/types.nim
new file mode 100644
index 0000000..b82c831
--- /dev/null
+++ b/tests/helpers/types.nim
@@ -0,0 +1,9 @@
+import swarmsim/engine/message
+
+type
+  FreelyTypedMessage* = ref object of Message
+    ## A `FreelyTypedMessage` is a `Message` that can be of any type.
+    ##
+    messageType*: string
+
+method `messageType`*(self: FreelyTypedMessage): string = self.messageType
diff --git a/tests/lib/twithtypeid.nim b/tests/lib/twithtypeid.nim
new file mode 100644
index 0000000..2f6b881
--- /dev/null
+++ b/tests/lib/twithtypeid.nim
@@ -0,0 +1,38 @@
+import unittest
+
+import swarmsim/lib/withtypeid
+
+withTypeId:
+  type
+    Foo = object of RootObj
+    Bar* = object of RootObj
+    Qux* = ref object of RootObj
+
+    FooBar = object of Bar
+
+type NonAnnotated = object of RootObj
+type NonAnnotatedRef = ref object of RootObj
+
+suite "withtypeid":
+  test "should allow querying a type for its id":
+    check(Foo.typeId == "Foo")
+    check(Bar.typeId == "Bar")
+    check(Qux.typeId == "Qux")
+
+  test "should allow querying an instance for its id":
+    check(Bar().typeId == "Bar")
+    check(Foo().typeId == "Foo")
+    check(Qux().typeId == "Qux")
+
+  test "should correctly return the id of the concrete type when upcasted":
+    let instance: Bar = FooBar()
+    check(instance.typeId == "FooBar")
+
+  test "should raise an error when trying to query the id of a non-annotated type":
+    expect(Defect):
+      discard NonAnnotated().typeId
+      # Note we don't need NonAnnotated.typeId as that won't even compile.
+
+  test "should raise an error when trying to query the id of a non-annotated ref type":
+    expect(Defect):
+      discard NonAnnotatedRef().typeId