Fix #3248: support product-seq patterns

liufengyun · liufengyun · commit 855adf82111f · 2019-03-05T14:22:01.000+01:00
diff --git a/compiler/src/dotty/tools/dotc/transform/PatternMatcher.scala b/compiler/src/dotty/tools/dotc/transform/PatternMatcher.scala
@@ -8,7 +8,7 @@ import collection.mutable
 import Symbols._, Contexts._, Types._, StdNames._, NameOps._
 import ast.Trees._
 import util.Spans._
-import typer.Applications.{isProductMatch, isGetMatch, productSelectors}
+import typer.Applications.{isProductMatch, isGetMatch, isProductSeqMatch, productSelectors, productArity}
 import SymUtils._
 import Flags._, Constants._
 import Decorators._
@@ -262,6 +262,8 @@ object PatternMatcher {
 
       /** Plan for matching the sequence in `getResult` against sequence elements
        *  and a possible last varargs argument `args`.
+       *
+       *  `getResult` could also be a product, where the last element is a sequence of elements.
        */
       def unapplySeqPlan(getResult: Symbol, args: List[Tree]): Plan = args.lastOption match {
         case Some(VarArgPattern(arg)) =>
@@ -286,6 +288,22 @@ object PatternMatcher {
           matchElemsPlan(getResult, args, exact = true, onSuccess)
       }
 
+      /** Plan for matching the sequence in `getResult` against sequence elements
+       *  and a possible last varargs argument `args`.
+       *
+       *  `getResult` is a product, where the last element is a sequence of elements.
+       */
+      def unapplyProductSeqPlan(getResult: Symbol, args: List[Tree], arity: Int): Plan = {
+        assert(arity <= args.size + 1)
+        val selectors = productSelectors(getResult.info).map(ref(getResult).select(_))
+
+        val matchSeq =
+          letAbstract(selectors.last) { seqResult =>
+            unapplySeqPlan(seqResult, args.drop(arity - 1))
+          }
+        matchArgsPlan(selectors.take(arity - 1), args.take(arity - 1), matchSeq)
+      }
+
       /** Plan for matching the result of an unapply against argument patterns `args` */
       def unapplyPlan(unapp: Tree, args: List[Tree]): Plan = {
         def caseClass = unapp.symbol.owner.linkedClass
@@ -306,18 +324,34 @@ object PatternMatcher {
                 .map(ref(unappResult).select(_))
               matchArgsPlan(selectors, args, onSuccess)
             }
+            else if (isProductSeqMatch(unapp.tpe.widen, args.length, unapp.sourcePos) && !isUnapplySeq) {
+              val arity = productArity(unapp.tpe.widen, unapp.sourcePos)
+              unapplyProductSeqPlan(unappResult, args, arity)
+            }
             else {
               assert(isGetMatch(unapp.tpe))
               val argsPlan = {
                 val get = ref(unappResult).select(nme.get, _.info.isParameterless)
+                val arity = productArity(get.tpe, unapp.sourcePos)
                 if (isUnapplySeq)
-                  letAbstract(get)(unapplySeqPlan(_, args))
+                  letAbstract(get) { getResult =>
+                    if (arity > 0) unapplyProductSeqPlan(getResult, args, arity)
+                    else unapplySeqPlan(getResult, args)
+                  }
                 else
                   letAbstract(get) { getResult =>
-                    val selectors =
-                      if (args.tail.isEmpty) ref(getResult) :: Nil
-                      else productSelectors(get.tpe).map(ref(getResult).select(_))
-                    matchArgsPlan(selectors, args, onSuccess)
+                    if (args.tail.isEmpty) // Single pattern takes precedence
+                      matchArgsPlan(ref(getResult) :: Nil, args, onSuccess)
+                    else if (isProductMatch(get.tpe, args.length, unapp.sourcePos)) {
+                      val sels = productSelectors(get.tpe).map(ref(getResult).select(_))
+                      matchArgsPlan(sels, args, onSuccess)
+                    }
+                    else if (isProductSeqMatch(get.tpe, args.length, unapp.sourcePos))
+                      unapplyProductSeqPlan(getResult, args, arity)
+                    else { // name-based
+                      val sels = productSelectors(get.tpe).map(ref(getResult).select(_))
+                      matchArgsPlan(sels, args, onSuccess)
+                    }
                   }
               }
               TestPlan(NonEmptyTest, unappResult, unapp.span, argsPlan)
diff --git a/compiler/src/dotty/tools/dotc/typer/Applications.scala b/compiler/src/dotty/tools/dotc/typer/Applications.scala
@@ -47,11 +47,22 @@ object Applications {
 
   /** Does `tp` fit the "product match" conditions as an unapply result type
    *  for a pattern with `numArgs` subpatterns?
-   *  This is the case of `tp` has members `_1` to `_N` where `N == numArgs`.
+   *  This is the case if `tp` has members `_1` to `_N` where `N == numArgs`.
    */
   def isProductMatch(tp: Type, numArgs: Int, errorPos: SourcePosition = NoSourcePosition)(implicit ctx: Context): Boolean =
     numArgs > 0 && productArity(tp, errorPos) == numArgs
 
+  /** Does `tp` fit the "product-seq match" conditions as an unapply result type
+   *  for a pattern with `numArgs` subpatterns?
+   *  This is the case if (1) `tp` has members `_1` to `_N` where `N <= numArgs + 1`.
+   *                      (2) `tp._N` conforms to Seq match
+   */
+  def isProductSeqMatch(tp: Type, numArgs: Int, errorPos: SourcePosition = NoSourcePosition)(implicit ctx: Context): Boolean = {
+    val arity = productArity(tp, errorPos)
+    arity > 0 && arity <= numArgs + 1 &&
+      unapplySeqTypeElemTp(productSelectorTypes(tp, errorPos).last).exists
+  }
+
   /** Does `tp` fit the "get match" conditions as an unapply result type?
    *  This is the case of `tp` has a `get` member as well as a
    *  parameterless `isEmpty` member of result type `Boolean`.
@@ -60,6 +71,39 @@ object Applications {
     extractorMemberType(tp, nme.isEmpty, errorPos).isRef(defn.BooleanClass) &&
     extractorMemberType(tp, nme.get, errorPos).exists
 
+  /** If `getType` is of the form:
+    *  ```
+    *  {
+    *    def lengthCompare(len: Int): Int // or, def length: Int
+    *    def apply(i: Int): T = a(i)
+    *    def drop(n: Int): scala.Seq[T]
+    *    def toSeq: scala.Seq[T]
+    *  }
+    *  ```
+    *  returns `T`, otherwise NoType.
+    */
+  def unapplySeqTypeElemTp(getTp: Type)(implicit ctx: Context): Type = {
+    def lengthTp = ExprType(defn.IntType)
+    def lengthCompareTp = MethodType(List(defn.IntType), defn.IntType)
+    def applyTp(elemTp: Type) = MethodType(List(defn.IntType), elemTp)
+    def dropTp(elemTp: Type) = MethodType(List(defn.IntType), defn.SeqType.appliedTo(elemTp))
+    def toSeqTp(elemTp: Type) = ExprType(defn.SeqType.appliedTo(elemTp))
+
+    // the result type of `def apply(i: Int): T`
+    val elemTp = getTp.member(nme.apply).suchThat(_.info <:< applyTp(WildcardType)).info.resultType
+
+    def hasMethod(name: Name, tp: Type) =
+      getTp.member(name).suchThat(getTp.memberInfo(_) <:< tp).exists
+
+    val isValid =
+      elemTp.exists &&
+      (hasMethod(nme.lengthCompare, lengthCompareTp) || hasMethod(nme.length, lengthTp)) &&
+      hasMethod(nme.drop, dropTp(elemTp)) &&
+      hasMethod(nme.toSeq, toSeqTp(elemTp))
+
+    if (isValid) elemTp else NoType
+  }
+
   def productSelectorTypes(tp: Type, errorPos: SourcePosition)(implicit ctx: Context): List[Type] = {
     def tupleSelectors(n: Int, tp: Type): List[Type] = {
       val sel = extractorMemberType(tp, nme.selectorName(n), errorPos)
@@ -89,9 +133,17 @@ object Applications {
     if (args.length > 1 && !(tp.derivesFrom(defn.SeqClass))) {
       val sels = productSelectorTypes(tp, pos)
       if (sels.length == args.length) sels
+      else if (isProductSeqMatch(tp, args.length, pos)) productSeqSelectors(tp, args, pos)
       else tp :: Nil
     } else tp :: Nil
 
+  def productSeqSelectors(tp: Type, args: List[untpd.Tree], pos: SourcePosition)(implicit ctx: Context): List[Type] = {
+      val selTps = productSelectorTypes(tp, pos)
+      val arity = selTps.length
+      val elemTp = unapplySeqTypeElemTp(selTps.last)
+      (0 until args.length).map(i => if (i < arity - 1) selTps(i) else elemTp).toList
+    }
+
   def unapplyArgs(unapplyResult: Type, unapplyFn: Tree, args: List[untpd.Tree], pos: SourcePosition)(implicit ctx: Context): List[Type] = {
 
     val unapplyName = unapplyFn.symbol.name
@@ -103,43 +155,11 @@ object Applications {
       Nil
     }
 
-    /** If `getType` is of the form:
-     *  ```
-     *  {
-     *    def lengthCompare(len: Int): Int // or, def length: Int
-     *    def apply(i: Int): T = a(i)
-     *    def drop(n: Int): scala.Seq[T]
-     *    def toSeq: scala.Seq[T]
-     *  }
-     *  ```
-     *  returns `T`, otherwise NoType.
-     */
-    def unapplySeqTypeElemTp(getTp: Type): Type = {
-      def lengthTp = ExprType(defn.IntType)
-      def lengthCompareTp = MethodType(List(defn.IntType), defn.IntType)
-      def applyTp(elemTp: Type) = MethodType(List(defn.IntType), elemTp)
-      def dropTp(elemTp: Type) = MethodType(List(defn.IntType), defn.SeqType.appliedTo(elemTp))
-      def toSeqTp(elemTp: Type) = defn.SeqType.appliedTo(elemTp)
-
-      // the result type of `def apply(i: Int): T`
-      val elemTp = getTp.member(nme.apply).suchThat(_.info <:< applyTp(WildcardType)).info.resultType
-
-      def hasMethod(name: Name, tp: Type) =
-        getTp.member(name).suchThat(getTp.memberInfo(_) <:< tp).exists
-
-      val isValid =
-        elemTp.exists &&
-        (hasMethod(nme.lengthCompare, lengthCompareTp) || hasMethod(nme.length, lengthTp)) &&
-        hasMethod(nme.drop, dropTp(elemTp)) &&
-        hasMethod(nme.toSeq, toSeqTp(elemTp))
-
-      if (isValid) elemTp else NoType
-    }
-
-    if (unapplyName == nme.unapplySeq) {
+    if (unapplyName == nme.unapplySeq) { //  && ctx.scala2Mode
       if (isGetMatch(unapplyResult, pos)) {
         val elemTp = unapplySeqTypeElemTp(getTp)
         if (elemTp.exists) args.map(Function.const(elemTp))
+        else if (isProductSeqMatch(getTp, args.length, pos)) productSeqSelectors(getTp, args, pos)
         else fail
       }
       else fail
@@ -148,6 +168,8 @@ object Applications {
       assert(unapplyName == nme.unapply)
       if (isProductMatch(unapplyResult, args.length, pos))
         productSelectorTypes(unapplyResult, pos)
+      else if (isProductSeqMatch(unapplyResult, args.length, pos))
+        productSeqSelectors(unapplyResult, args, pos)
       else if (isGetMatch(unapplyResult, pos))
         getUnapplySelectors(getTp, args, pos)
       else if (unapplyResult.widenSingleton isRef defn.BooleanClass)
diff --git a/tests/run/i3248.scala b/tests/run/i3248.scala
@@ -1,23 +1,27 @@
-object Test extends App {
+object Test {
   class Foo(val name: String, val children: Int *)
   object Foo {
-    def unapply(f: Foo) = Some((f.name, f.children))
+    def unapplySeq(f: Foo) = Some((f.name, f.children))
   }
 
-  def foo(f: Foo) = (f: Any) match {
-    case Foo(name, ns: _*) => ns.length
-    case List(ns: _*) => ns.length
+  def foo(f: Foo) = f match {
+    case Foo(name, ns : _*) =>
+      assert(name == "hello")
+      assert(ns(0) == 3)
+      assert(ns(1) == 5)
   }
 
-  case class Bar(val children: Int*)
-
-  def bar(f: Any) = f match {
-    case Bar(1, 2, 3) => 0
-    case Bar(a, b) => a + b
-    case Bar(ns: _*) => ns.length
+  def bar(f: Foo) = f match {
+    case Foo(name, x, y, ns : _*) =>
+      assert(name == "hello")
+      assert(x == 3)
+      assert(y == 5)
+      assert(ns.isEmpty)
   }
 
-  assert(bar(new Bar(1, 2, 3)) == 0)
-  assert(bar(new Bar(3, 2, 1)) == 3)
-  assert(foo(new Foo("name", 1, 2, 3)) == 3)
+  def main(args: Array[String]): Unit = {
+    val f = new Foo("hello", 3, 5)
+    foo(f)
+    bar(f)
+  }
 }
