<NME> TypeReplacement.kt
<BEF> package org.elm.lang.core.types

import com.intellij.psi.PsiElement
import org.elm.lang.core.diagnostics.ElmDiagnostic
import org.elm.lang.core.diagnostics.RecordBaseTypeError
import org.elm.lang.core.diagnostics.TypeArgumentCountError

/**
 * This class performs deep replacement of a set of [TyVar]s in a [Ty] with a set of new types,
 * which could also be [TyVar]s.
 *
 * It relies on the fact that [TyVar]s can be compared by identity. Vars in different scopes must
 * compare unequal, even if they have the same name.
 */
class TypeReplacement private constructor(
        // A map of variables that should be replaced to a pair of the ty to replace them with and
        // the psi element for the argument, which is used to show errors.
        private val replacements: Map<TyVar, Pair<PsiElement, Ty>>
) {
    companion object {
        /**
         * Replace types in a ty inferred from a type ref with arguments.
         *
         * @param element the reference element that [ty] was inferred from
         * @param ty the type to perform replacement in
         * @param paramTys the parameters of [ty] that will be replaced
         * @param argTys the arguments that will replaces the parameters.
         * @param argElements the PsiElements that [argTys] were inferred from. Must be the same size as [argTys].
         */
        fun replaceCall(
                element: PsiElement,
                ty: Ty,
                paramTys: List<TyVar>,
                argTys: List<Ty>,
                argElements: List<PsiElement>
        ): ParameterizedInferenceResult<Ty> {
            require(argTys.size == argElements.size) { "mismatched arg sizes ${argTys.size} != ${argElements.size}" }

            if (paramTys.size != argTys.size) {
                val error = TypeArgumentCountError(element, argTys.size, paramTys.size)
                return ParameterizedInferenceResult(listOf(error), TyUnknown())
            }

            if (paramTys.isEmpty()) {
                return ParameterizedInferenceResult(emptyList(), ty)
            }

            val replacements = paramTys.indices.associate { i -> paramTys[i] to (argElements[i] to argTys[i]) }
            val typeReplacement = TypeReplacement(replacements)
            val newTy = typeReplacement.replace(ty)
            return ParameterizedInferenceResult(typeReplacement.diagnostics, newTy)
        }
    }

    private val diagnostics = mutableListOf<ElmDiagnostic>()

    private fun replace(ty: Ty): Ty = when (ty) {
        is TyVar -> replacements[ty]?.second ?: ty
        is TyTuple -> TyTuple(ty.types.map { replace(it) }, replace(ty.alias))
        is TyFunction -> TyFunction(ty.parameters.map { replace(it) }, replace(ty.ret), replace(ty.alias))
        is TyUnknown -> TyUnknown(replace(ty.alias))
        is TyUnion -> replaceUnion(ty)
        is TyRecord -> replaceRecord(ty)
         * Make all type variables in a [ty] flexible.
         *
         * Type variables in function annotations are rigid when bound to parameters; in all other
         * cases vars are flexible. This function is used to make vars inferred as rigid into flexible
         * when calling functions.
         */
        fun flexify(ty: Ty): Ty {
            return TypeReplacement(emptyMap(), freshen = false, varsToRemainRigid = emptyList(), keepRecordsMutable = false).replace(ty)
        }

        /**
         * Freeze any fields references present in records in a [ty].
         *
         * Use this to prevent cached values from being altered. Note that this is an in-place
         * change rather than a copy.
         */
        fun freeze(ty: Ty) {
            when (ty) {
                is TyVar -> Unit
    }

    private fun replaceRecord(ty: TyRecord): Ty {
        val baseTy = replacements[ty.baseTy]?.second
        val baseFields = (baseTy as? TyRecord)?.fields.orEmpty()

        if (baseTy != null && isInferable(baseTy) && baseTy !is TyRecord) {
            diagnostics += RecordBaseTypeError(replacements[ty.baseTy]!!.first, baseTy)
        }

        val declaredFields = ty.fields.mapValues { (_, it) -> replace(it) }

        val newBaseTy = when(baseTy) {
            // If the base ty of the argument is a record, use it's base ty, which might be null.
            is TyRecord -> baseTy.baseTy
            // If it's another variable, use it as-is
            else -> baseTy
        }
        return TyRecord(baseFields + declaredFields, newBaseTy, replace(ty.alias))
    }
}

        // If we wouldn't change anything, return the original ty to avoid duplicating the object
        if (!wouldChange(ty)) return ty
        return when (ty) {
            is TyVar -> getReplacement(ty) ?: ty
            is TyTuple -> replaceTuple(ty)
            is TyFunction -> replaceFunction(ty)
            is TyUnknown -> TyUnknown(replace(ty.alias))
            is TyUnion -> replaceUnion(ty)
            is TyRecord -> replaceRecord(ty.fields, ty.baseTy, ty.alias, ty.fieldReferences, wasMutable = false)
            is TyUnit, TyInProgressBinding -> ty
            is MutableTyRecord -> replaceRecord(ty.fields, ty.baseTy, null, ty.fieldReferences, wasMutable = true)
        }
    }

    /*
     * Although aliases are not used in ty comparisons, we still need to do replacement on them to
     * render their call sites correctly.
     * e.g.
     *     type alias A a = ...
     *     main : A ()
     * We replace the parameter of the AliasInfo in the return value of the main function with
     * TyUnit so that it renders as `A ()` rather than `A a`.
     */
    private fun replace(aliasInfo: AliasInfo?) = aliasInfo?.let { info ->
        info.copy(parameters = info.parameters.map { replace(it) }.optimizeReadOnlyList())
    }

    private fun replaceTuple(ty: TyTuple): TyTuple {
        return TyTuple(ty.types.map { replace(it) }.optimizeReadOnlyList(), replace(ty.alias))
    }

    private fun replaceFunction(ty: TyFunction): TyFunction {
        val parameters = ty.parameters.map { replace(it) }.optimizeReadOnlyList()
        return TyFunction(parameters, replace(ty.ret), replace(ty.alias)).uncurry()
    }

    private fun replaceUnion(ty: TyUnion): TyUnion {
        // fast path for common cases like Basics.Int
        if (ty.parameters.isEmpty() && ty.alias == null) return ty

        val parameters = ty.parameters.map { replace(it) }.optimizeReadOnlyList()
        return TyUnion(ty.module, ty.name, parameters, replace(ty.alias))
    }

    private fun replaceRecord(
            fields: Map<String, Ty>,
            baseTy: Ty?,
            alias: AliasInfo?,
            fieldReferences: RecordFieldReferenceTable,
            wasMutable: Boolean
    ): Ty {
        val oldBase = if (baseTy == null || baseTy !is TyVar) null else getReplacement(baseTy)
        val newBase = when (oldBase) {
            // If the base ty of the argument is a record, use its base ty, which might be null.
            is TyRecord -> oldBase.baseTy
            // If it wasn't substituted, leave it as-is
            null -> when {
                // Although if it's a record, we might need to make it immutable
                baseTy is MutableTyRecord && !keepRecordsMutable -> replace(baseTy)
                else -> baseTy
            }
            // If it's another variable, use it
            else -> oldBase
        }

        val baseFields = (oldBase as? TyRecord)?.fields.orEmpty()
        val baseFieldRefs = (oldBase as? TyRecord)?.fieldReferences

        // Make the new map as small as we can for these fields. HashMap always uses a power of two
        // element array, so we can't avoid all wasted memory.
        val initialCapacity = baseFields.size + fields.keys.count { it !in baseFields }
        val newFields = LinkedHashMap<String, Ty>(initialCapacity, 1f)

        // Don't use putAll here, since that function will resize the table to hold (size + 1) elements
        baseFields.forEach { (k, v) -> newFields[k] = v }
        fields.mapValuesTo(newFields) { (_, it) -> replace(it) }

        val newFieldReferences = when {
            baseFieldRefs == null || baseFieldRefs.isEmpty() -> fieldReferences
            fieldReferences.frozen -> fieldReferences + baseFieldRefs
            else -> {
                // The new record shares its references table with the old record. That allows us to track
                // references back to expressions inside nested declarations even when the record has been
                // freshened or replaced.
                fieldReferences.apply { addAll(baseFieldRefs) }
            }
        }

        return if (wasMutable && keepRecordsMutable) MutableTyRecord(newFields, newBase, newFieldReferences)
        else TyRecord(newFields, newBase, replace(alias), newFieldReferences)
    }

    // When we replace a var, the new ty may itself contain vars, and so we need to recursively
    // replace the ty before we can replace the var we were given as an argument.
    // After the recursive replacement, we avoid repeating work by storing the final ty and tracking
    // of the fact that its replacement is complete with the `hasBeenAccessed` flag.
    private fun getReplacement(key: TyVar): Ty? {
        val (hasBeenAccessed, storedTy) = replacements[key] ?: return when {
            freshen || varsToRemainRigid != null -> {
                if (key.rigid && (varsToRemainRigid == null || key in varsToRemainRigid)) null // never freshen rigid vars
                else TyVar(key.name, rigid = false).also { replacements[key] = true to it }
            }
            else -> null
        }

        if (hasBeenAccessed) return storedTy

        val replacedVal = replace(storedTy)
        replacements[key] = true to replacedVal
        return replacedVal
    }
}

// Copied from the stdlib, which doesn't apply this for `map` and `filter`
private fun <T> List<T>.optimizeReadOnlyList() = when (size) {
    0 -> emptyList()
    1 -> listOf(this[0])
    else -> this
}

<MSG> Unify type variables in function calls (#254)


<DFF> @@ -1,10 +1,5 @@
 package org.elm.lang.core.types
 
-import com.intellij.psi.PsiElement
-import org.elm.lang.core.diagnostics.ElmDiagnostic
-import org.elm.lang.core.diagnostics.RecordBaseTypeError
-import org.elm.lang.core.diagnostics.TypeArgumentCountError
-
 /**
  * This class performs deep replacement of a set of [TyVar]s in a [Ty] with a set of new types,
  * which could also be [TyVar]s.
@@ -12,52 +7,60 @@ import org.elm.lang.core.diagnostics.TypeArgumentCountError
  * It relies on the fact that [TyVar]s can be compared by identity. Vars in different scopes must
  * compare unequal, even if they have the same name.
  */
-class TypeReplacement private constructor(
-        // A map of variables that should be replaced to a pair of the ty to replace them with and
-        // the psi element for the argument, which is used to show errors.
-        private val replacements: Map<TyVar, Pair<PsiElement, Ty>>
+class TypeReplacement(
+        // A map of variables that should be replaced to the ty to replace them with
+        private val replacements: Map<TyVar, Ty>
 ) {
     companion object {
         /**
-         * Replace types in a ty inferred from a type ref with arguments.
-         *
-         * @param element the reference element that [ty] was inferred from
-         * @param ty the type to perform replacement in
-         * @param paramTys the parameters of [ty] that will be replaced
-         * @param argTys the arguments that will replaces the parameters.
-         * @param argElements the PsiElements that [argTys] were inferred from. Must be the same size as [argTys].
+         * Replace vars in [ty] according to [replacements].
          */
-        fun replaceCall(
-                element: PsiElement,
-                ty: Ty,
-                paramTys: List<TyVar>,
-                argTys: List<Ty>,
-                argElements: List<PsiElement>
-        ): ParameterizedInferenceResult<Ty> {
-            require(argTys.size == argElements.size) { "mismatched arg sizes ${argTys.size} != ${argElements.size}" }
-
-            if (paramTys.size != argTys.size) {
-                val error = TypeArgumentCountError(element, argTys.size, paramTys.size)
-                return ParameterizedInferenceResult(listOf(error), TyUnknown())
-            }
+        fun replace(ty: Ty, replacements: Map<TyVar, Ty>): Ty {
+            if (replacements.isEmpty()) return ty
+            return TypeReplacement(replacements).replace(ty)
+        }
 
-            if (paramTys.isEmpty()) {
-                return ParameterizedInferenceResult(emptyList(), ty)
+        /**
+         * Replace vars in [ty] according to [replacements].
+         *
+         * If the values of [replacements] can contain [TyVar]s that occur in the keys, use this
+         * rather than [replace].
+         *
+         * This will perform replacement several times until there is nothing left to replace. Then
+         * all remaining [TyVar]s in the type are replaced with new vars with the same name to
+         * enforce scoping.
+         */
+        fun deepReplace(ty: Ty, replacements: Map<TyVar, Ty>): Ty {
+            if (replacements.isEmpty()) return ty
+            val tr = TypeReplacement(replacements)
+            var newTy = ty
+            repeat(5) {
+                val next = tr.replace(newTy)
+                if (next == newTy) {
+                    return freshenVars(newTy)
+                }
+                newTy = next
             }
+            // exceeded the recursion threshold, don't replace anything
+            return ty
+        }
 
-            val replacements = paramTys.indices.associate { i -> paramTys[i] to (argElements[i] to argTys[i]) }
-            val typeReplacement = TypeReplacement(replacements)
-            val newTy = typeReplacement.replace(ty)
-            return ParameterizedInferenceResult(typeReplacement.diagnostics, newTy)
+        private fun freshenVars(ty: Ty): Ty {
+            return TypeReplacement(MutableMapWithDefault { TyVar(it.name) }).replace(ty)
         }
     }
 
-    private val diagnostics = mutableListOf<ElmDiagnostic>()
-
     private fun replace(ty: Ty): Ty = when (ty) {
-        is TyVar -> replacements[ty]?.second ?: ty
+        is TyVar -> replacements[ty] ?: ty
         is TyTuple -> TyTuple(ty.types.map { replace(it) }, replace(ty.alias))
-        is TyFunction -> TyFunction(ty.parameters.map { replace(it) }, replace(ty.ret), replace(ty.alias))
+        is TyFunction -> {
+            val parameters = ty.parameters.map {
+                replace(it)
+            }
+            val ret = replace(ty.ret)
+            val alias = replace(ty.alias)
+            TyFunction(parameters, ret, alias).uncurry()
+        }
         is TyUnknown -> TyUnknown(replace(ty.alias))
         is TyUnion -> replaceUnion(ty)
         is TyRecord -> replaceRecord(ty)
@@ -83,22 +86,42 @@ class TypeReplacement private constructor(
     }
 
     private fun replaceRecord(ty: TyRecord): Ty {
-        val baseTy = replacements[ty.baseTy]?.second
-        val baseFields = (baseTy as? TyRecord)?.fields.orEmpty()
-
-        if (baseTy != null && isInferable(baseTy) && baseTy !is TyRecord) {
-            diagnostics += RecordBaseTypeError(replacements[ty.baseTy]!!.first, baseTy)
+        val replacedBase = if (ty.baseTy == null) null else replacements[ty.baseTy]
+        val newBaseTy = when (replacedBase) {
+            // If the base ty of the argument is a record, use it's base ty, which might be null.
+            is TyRecord -> replacedBase.baseTy
+            // If it wasn't substituted, leave it as-is
+            null -> ty.baseTy
+            // If it's another variable, use it
+            else -> replacedBase
         }
 
         val declaredFields = ty.fields.mapValues { (_, it) -> replace(it) }
+        val baseFields = (replacedBase as? TyRecord)?.fields.orEmpty()
 
-        val newBaseTy = when(baseTy) {
-            // If the base ty of the argument is a record, use it's base ty, which might be null.
-            is TyRecord -> baseTy.baseTy
-            // If it's another variable, use it as-is
-            else -> baseTy
-        }
         return TyRecord(baseFields + declaredFields, newBaseTy, replace(ty.alias))
     }
 }
 
+
+private class MutableMapWithDefault<K, V>(
+        val map: MutableMap<K, V> = mutableMapOf(),
+        private val default: (key: K) -> V
+) : MutableMap<K, V> {
+    override fun equals(other: Any?): Boolean = map.equals(other)
+    override fun hashCode(): Int = map.hashCode()
+    override fun toString(): String = map.toString()
+    override val size: Int get() = map.size
+    override fun isEmpty(): Boolean = map.isEmpty()
+    override fun containsKey(key: K): Boolean = map.containsKey(key)
+    override fun containsValue(value: @UnsafeVariance V): Boolean = map.containsValue(value)
+    override fun get(key: K): V? = map.getOrPut(key) { default(key) }
+    override val keys: MutableSet<K> get() = map.keys
+    override val values: MutableCollection<V> get() = map.values
+    override val entries: MutableSet<MutableMap.MutableEntry<K, V>> get() = map.entries
+
+    override fun put(key: K, value: V): V? = map.put(key, value)
+    override fun remove(key: K): V? = map.remove(key)
+    override fun putAll(from: Map<out K, V>) = map.putAll(from)
+    override fun clear() = map.clear()
+}

Unify type variables in function calls (#254)