FIR: Support expansion of flexible type-alias-based types in inference

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/types/ConeTypeContext.kt

@@ -398,6 +398,15 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext, Ty
     override fun prepareType(type: KotlinTypeMarker): KotlinTypeMarker {
         return when (type) {
             is ConeClassLikeType -> type.fullyExpandedType(session)
+            is ConeFlexibleType -> {
+                val lowerBound = prepareType(type.lowerBound)
+                if (lowerBound === type.lowerBound) return type
+
+                ConeFlexibleType(
+                    lowerBound as ConeKotlinType,
+                    prepareType(type.upperBound) as ConeKotlinType
+                )
+            }
             else -> type
         }
     }

compiler/fir/resolve/testData/resolveWithStdlib/j+k/flexibleTypeAliases.kt

@@ -0,0 +1,35 @@
+// FULL_JDK
+// FILE: imm/Map.java
+package imm;
+public interface Map<K, V> {
+    Option<V> get(k K);
+}
+
+// FILE: imm/Set.java
+package imm;
+public interface Set<E> {
+}
+
+// FILE: imm/Option.java
+package imm;
+public interface Option<T> {
+    T getOrElse(T other);
+}
+
+// FILE: imm/LinkedHashSet.java
+package imm;
+public class LinkedHashSet<E> extends Set<E> {
+    public static <T> LinkedHashSet<T> empty() { return null; }
+}
+
+// FILE: main.kt
+
+typealias ImmutableMap<K, V> = imm.Map<K, V>
+typealias ImmutableSet<E> = imm.Set<E>
+typealias ImmutableLinkedHashSet<E> = imm.LinkedHashSet<E>
+
+private typealias ImmutableMultimap<K, V> = ImmutableMap<K, ImmutableSet<V>>
+
+private fun <K, V> ImmutableMultimap<K, V>.put(key: K, value: V) {
+    this[key].getOrElse(ImmutableLinkedHashSet.empty<V>())
+}

compiler/fir/resolve/testData/resolveWithStdlib/j+k/flexibleTypeAliases.txt

@@ -0,0 +1,8 @@
+FILE: main.kt
+    public final typealias ImmutableMap<K, V> = R|imm/Map<K, V>|
+    public final typealias ImmutableSet<E> = R|imm/Set<E>|
+    public final typealias ImmutableLinkedHashSet<E> = R|imm/LinkedHashSet<E>|
+    private final typealias ImmutableMultimap<K, V> = R|ImmutableMap<K, ImmutableSet<V>>|
+    private final fun <K, V> R|ImmutableMultimap<K, V>|.put(key: R|K|, value: R|V|): R|kotlin/Unit| {
+        this@R|/put|.R|FakeOverride<imm/Map.get: R|ft<imm/Option<ft<ImmutableSet<V>, ImmutableSet<V>?>!>, imm/Option<ft<ImmutableSet<V>, ImmutableSet<V>?>!>?>!|>|(R|<local>/key|).R|FakeOverride<imm/Option.getOrElse: R|ft<ImmutableSet<V>, ImmutableSet<V>?>!|>|(Q|ImmutableLinkedHashSet|.R|imm/LinkedHashSet.empty|<R|V|>())
+    }

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsWithStdlibTestGenerated.java

@@ -585,6 +585,11 @@ public class FirDiagnosticsWithStdlibTestGenerated extends AbstractFirDiagnostic
             runTest("compiler/fir/resolve/testData/resolveWithStdlib/j+k/FieldSubstitution.kt");
         }
 
+        @TestMetadata("flexibleTypeAliases.kt")
+        public void testFlexibleTypeAliases() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolveWithStdlib/j+k/flexibleTypeAliases.kt");
+        }
+
         @TestMetadata("FunctionTypeInJava.kt")
         public void testFunctionTypeInJava() throws Exception {
             runTest("compiler/fir/resolve/testData/resolveWithStdlib/j+k/FunctionTypeInJava.kt");