Add userType information to Field and Message handles.

This information is needed in order to perform modifications where handles
play crucial role of pointing the fragment of a tree that should be
modified.

Signed-off-by: Grzegorz Kossakowski <grek@google.com>

3 files changed, 11 insertions, 7 deletions

diff --git a/src/main/scala/com/google/gimd/query/Handle.scala b/src/main/scala/com/google/gimd/query/Handle.scala
index aeb35dc..de42d5c 100644
--- a/src/main/scala/com/google/gimd/query/Handle.scala
+++ b/src/main/scala/com/google/gimd/query/Handle.scala
@@ -27,7 +27,9 @@ abstract case class RootHandle() extends Handle
 
 /** Handle to a message by itself. */
 final case class MessageHandle(
-  message: Message
+  message: Message,
+  /** userType for Message that this handle is pointing at */
+  userType: UserType[_]
 ) extends InnerHandle
 
 /** Handle to a message stored within a field. */
@@ -37,7 +39,9 @@ final case class FieldHandle(
   // parent, rather than using the object.
   //
   field: MessageField,
-  child: InnerHandle
+  child: InnerHandle,
+  /** userType defined for message which contains given field */
+  userType: UserType[_]
 ) extends InnerHandle
 
 /** Handle to a top level message. */
diff --git a/src/main/scala/com/google/gimd/query/MessageQuery.scala b/src/main/scala/com/google/gimd/query/MessageQuery.scala
index f88b0da..9dfac04 100644
--- a/src/main/scala/com/google/gimd/query/MessageQuery.scala
+++ b/src/main/scala/com/google/gimd/query/MessageQuery.scala
@@ -29,7 +29,7 @@ object MessageQuery {
     if (p.isType(ut.userTypeClass)) {
       val obj = ut.toUserObject(m).asInstanceOf[U]
       if (p.isMatch(obj))
-        Iterator.single( (MessageHandle(m), obj) )
+        Iterator.single( (MessageHandle(m, ut), obj) )
       else
         Iterator.empty
     } else
@@ -45,5 +45,5 @@ object MessageQuery {
       if field.isInstanceOf[MessageField]
       f = field.asInstanceOf[MessageField]
       (handle, userObject) <- simpleQuery(member.userType, f.value, p)
-    } yield (FieldHandle(f, handle), userObject)
+    } yield (FieldHandle(f, handle, ut), userObject)
 }
diff --git a/src/test/scala/com/google/gimd/query/MessageQueryTestCase.scala b/src/test/scala/com/google/gimd/query/MessageQueryTestCase.scala
index 3db5968..df1c5d6 100644
--- a/src/test/scala/com/google/gimd/query/MessageQueryTestCase.scala
+++ b/src/test/scala/com/google/gimd/query/MessageQueryTestCase.scala
@@ -85,13 +85,13 @@ class MessageQueryTestCase {
     MessageQuery.simpleQuery(TreeNodeType, message, p)
   }
 
-  private def handle(message: Message): InnerHandle = MessageHandle(message)
+  private def handle(message: Message): InnerHandle = MessageHandle(message, TreeNodeType)
   private def handle(p: (String, Message)*): InnerHandle = {
     val (name, message) = p(0)
     val field = MessageField(name, message)
     if (p.size == 1) {
-      FieldHandle(field, handle(message))
+      FieldHandle(field, handle(message), TreeNodeType)
     } else
-      FieldHandle(field, handle(p.drop(1): _*))
+      FieldHandle(field, handle(p.drop(1): _*), TreeNodeType)
   }
 }