Handle hk lambdas in tasty

3 files changed, 41 insertions, 32 deletions

diff --git a/compiler/src/dotty/tools/dotc/core/tasty/TastyFormat.scala b/compiler/src/dotty/tools/dotc/core/tasty/TastyFormat.scala
index 98577f3c5..8b2255e94 100644
--- a/compiler/src/dotty/tools/dotc/core/tasty/TastyFormat.scala
+++ b/compiler/src/dotty/tools/dotc/core/tasty/TastyFormat.scala
@@ -151,8 +151,9 @@ Standard-Section: "ASTs" TopLevelStat*
                   BIND           Length boundName_NameRef bounds_Type
                                         // for type-variables defined in a type pattern
                   BYNAMEtype            underlying_Type
-                  POLYtype       Length result_Type NamesTypes      // variance encoded in front of name: +/-/(nothing)
+                  POLYtype       Length result_Type NamesTypes
                   METHODtype     Length result_Type NamesTypes      // needed for refinements
+                  TYPELAMBDAtype Length result_Type NamesTypes      // variance encoded in front of name: +/-/(nothing)
                   PARAMtype      Length binder_ASTref paramNum_Nat  // needed for refinements
                   SHARED                type_ASTRef
   NamesTypes    = NameType*
@@ -345,9 +346,10 @@ object TastyFormat {
   final val ORtpt = 169
   final val METHODtype = 170
   final val POLYtype = 171
-  final val POLYtpt = 172
-  final val PARAMtype = 173
-  final val ANNOTATION = 174
+  final val TYPELAMBDAtype = 172
+  final val LAMBDAtpt = 173
+  final val PARAMtype = 174
+  final val ANNOTATION = 175
 
   final val firstSimpleTreeTag = UNITconst
   final val firstNatTreeTag = SHARED
@@ -397,7 +399,7 @@ object TastyFormat {
        | SINGLETONtpt
        | REFINEDtpt
        | APPLIEDtpt
-       | POLYtpt
+       | LAMBDAtpt
        | TYPEBOUNDStpt
        | ANNOTATEDtpt
        | ANDtpt
@@ -528,8 +530,9 @@ object TastyFormat {
     case BYNAMEtype => "BYNAMEtype"
     case BYNAMEtpt => "BYNAMEtpt"
     case POLYtype => "POLYtype"
-    case POLYtpt => "POLYtpt"
     case METHODtype => "METHODtype"
+    case TYPELAMBDAtype => "TYPELAMBDAtype"
+    case LAMBDAtpt => "LAMBDAtpt"
     case PARAMtype => "PARAMtype"
     case ANNOTATION => "ANNOTATION"
     case PRIVATEqualified => "PRIVATEqualified"
@@ -543,7 +546,7 @@ object TastyFormat {
     case VALDEF | DEFDEF | TYPEDEF | TYPEPARAM | PARAM | NAMEDARG | RETURN | BIND |
          SELFDEF | REFINEDtype => 1
     case RENAMED | PARAMtype => 2
-    case POLYtype | METHODtype => -1
+    case POLYtype | METHODtype | TYPELAMBDAtype => -1
     case _ => 0
   }
 }
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/TreePickler.scala b/compiler/src/dotty/tools/dotc/core/tasty/TreePickler.scala
index 4c600564d..fc9d9c447 100644
--- a/compiler/src/dotty/tools/dotc/core/tasty/TreePickler.scala
+++ b/compiler/src/dotty/tools/dotc/core/tasty/TreePickler.scala
@@ -253,12 +253,12 @@ class TreePickler(pickler: TastyPickler) {
     case tpe: ExprType =>
       writeByte(BYNAMEtype)
       pickleType(tpe.underlying)
-    case tpe: TypeLambda =>
-      writeByte(POLYtype)
-      pickleMethodic(tpe.resultType, tpe.paramNames, tpe.paramInfos)
+    case tpe: HKTypeLambda =>
+      pickleMethodic(TYPELAMBDAtype, tpe)
+    case tpe: PolyType /*if richTypes*/ => //###
+      pickleMethodic(POLYtype, tpe)
     case tpe: MethodType if richTypes =>
-      writeByte(METHODtype)
-      pickleMethodic(tpe.resultType, tpe.paramNames, tpe.paramInfos)
+      pickleMethodic(METHODtype, tpe)
     case tpe: TypeParamRef =>
       if (!pickleParamRef(tpe))
       // TODO figure out why this case arises in e.g. pickling AbstractFileReader.
@@ -281,13 +281,15 @@ class TreePickler(pickler: TastyPickler) {
     pickleName(qualifiedName(pkg))
   }
 
-  def pickleMethodic(result: Type, names: List[Name], types: List[Type])(implicit ctx: Context) =
+  def pickleMethodic(tag: Int, tpe: LambdaType)(implicit ctx: Context) = {
+    writeByte(tag)
     withLength {
-      pickleType(result, richTypes = true)
-      (names, types).zipped.foreach { (name, tpe) =>
+      pickleType(tpe.resultType, richTypes = true)
+      (tpe.paramNames, tpe.paramInfos).zipped.foreach { (name, tpe) =>
         pickleName(name); pickleType(tpe)
       }
     }
+  }
 
   def pickleParamRef(tpe: ParamRef)(implicit ctx: Context): Boolean = {
     val binder = pickledTypes.get(tpe.binder)
@@ -554,7 +556,7 @@ class TreePickler(pickler: TastyPickler) {
           writeByte(ANNOTATEDtpt)
           withLength { pickleTree(tree); pickleTree(annot.tree) }
         case LambdaTypeTree(tparams, body) =>
-          writeByte(POLYtpt)
+          writeByte(LAMBDAtpt)
           withLength { pickleParams(tparams); pickleTree(body) }
         case TypeBoundsTree(lo, hi) =>
           writeByte(TYPEBOUNDStpt)
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala b/compiler/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala
index e751e99fe..b83045a11 100644
--- a/compiler/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala
+++ b/compiler/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala
@@ -234,6 +234,19 @@ class TreeUnpickler(reader: TastyReader, tastyName: TastyName.Table, posUnpickle
           (nameReader.readParamNames(end), paramReader)
         }
 
+        def readMethodic[N <: Name, PInfo <: Type, LT <: LambdaType]
+            (companion: LambdaTypeCompanion[N, PInfo, LT], nameMap: Name => N): LT = {
+          val nameReader = fork
+          nameReader.skipTree() // skip result
+          val paramReader = nameReader.fork
+          val paramNames = nameReader.readParamNames(end).map(nameMap)
+          val result = companion(paramNames)(
+                pt => registeringType(pt, paramReader.readParamTypes[PInfo](end)),
+                pt => readType())
+          goto(end)
+          result
+        }
+
         val result =
           (tag: @switch) match {
             case SUPERtype =>
@@ -268,23 +281,14 @@ class TreeUnpickler(reader: TastyReader, tastyName: TastyName.Table, posUnpickle
               registerSym(start, sym)
               TypeRef.withFixedSym(NoPrefix, sym.name, sym)
             case POLYtype =>
-              val (paramNames, paramReader) = readNamesSkipParams
-              val result = PolyType(paramNames.map(_.toTypeName))(
-                pt => registeringType(pt, paramReader.readParamTypes[TypeBounds](end)),
-                pt => readType())
-              goto(end)
-              result
+              readMethodic(PolyType, _.toTypeName)
             case METHODtype =>
-              val (names, paramReader) = readNamesSkipParams
-              val result = MethodType(names.map(_.toTermName))(
-                mt => registeringType(mt, paramReader.readParamTypes[Type](end)),
-                mt => readType())
-              goto(end)
-              result
+              readMethodic(MethodType, _.toTermName)
+            case TYPELAMBDAtype =>
+              readMethodic(HKTypeLambda, _.toTypeName)
             case PARAMtype =>
               readTypeRef() match {
-                case binder: TypeLambda => binder.newParamRef(readNat())
-                case binder: MethodType => binder.newParamRef(readNat())
+                case binder: LambdaType => binder.newParamRef(readNat())
               }
             case CLASSconst =>
               ConstantType(Constant(readType()))
@@ -410,7 +414,7 @@ class TreeUnpickler(reader: TastyReader, tastyName: TastyName.Table, posUnpickle
     }
 
     def isAbstractType(ttag: Int)(implicit ctx: Context): Boolean = nextUnsharedTag match {
-      case POLYtpt =>
+      case LAMBDAtpt =>
         val rdr = fork
         rdr.reader.readByte()  // tag
         rdr.reader.readNat()   // length
@@ -1033,7 +1037,7 @@ class TreeUnpickler(reader: TastyReader, tastyName: TastyName.Table, posUnpickle
               OrTypeTree(readTpt(), readTpt())
             case ANNOTATEDtpt =>
               Annotated(readTpt(), readTerm())
-            case POLYtpt =>
+            case LAMBDAtpt =>
               val localCtx = localNonClassCtx
               val tparams = readParams[TypeDef](TYPEPARAM)(localCtx)
               val body = readTpt()(localCtx)