Added locator for untyped trees.

1 files changed, 13 insertions, 0 deletions

diff --git a/src/dotty/tools/dotc/core/UntypedTrees.scala b/src/dotty/tools/dotc/core/UntypedTrees.scala
index 82a160ab0..8cd857f95 100644
--- a/src/dotty/tools/dotc/core/UntypedTrees.scala
+++ b/src/dotty/tools/dotc/core/UntypedTrees.scala
@@ -74,5 +74,18 @@ object UntypedTrees {
 
   def ugen(implicit ctx: Context) =
     new UGen
+
+  implicit class UntypedTreeDecorator(val self: Tree) extends AnyVal {
+    def locateEnclosing(base: List[Tree], pos: Position): List[Tree] = {
+      def encloses(elem: Any) = elem match {
+        case t: Tree => t.envelope contains pos
+        case _ => false
+      }
+      base.productIterator find encloses match {
+        case Some(tree: Tree) => locateEnclosing(tree :: base, pos)
+        case none => base
+      }
+    }
+  }
 }