Introduce structured hash key

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/api/types/semtype/CacheableTypeDescriptor.java

@@ -11,15 +11,6 @@
  */
 public interface CacheableTypeDescriptor extends Type {
 
-    /**
-     * Check whether the type check result of this type descriptor should be cached. Can be used to avoid caching in
-     * cases where either directly doing the type check is cheaper or we can't determine if two instances of a type
-     * descriptor are equal without doing a type check.
-     *
-     * @return true if the type check result should be cached, false otherwise
-     */
-    boolean shouldCache();
-
     /**
      * Check whether the type check result of this type descriptor is cached for the given type descriptor.
      *
@@ -30,12 +21,12 @@ public interface CacheableTypeDescriptor extends Type {
     Optional<Boolean> cachedTypeCheckResult(Context cx, CacheableTypeDescriptor other);
 
     /**
-     * Cache the type check result of this type descriptor for the given type descriptor. Note that implementations of
-     * this method could choose to not cache the result if {@link #shouldCache()} returns false. In such cases, even
-     * after calling this method, {@link #cachedTypeCheckResult(Context, CacheableTypeDescriptor)} could return empty.
+     * Cache the type check result of this type descriptor for the given type descriptor.
      *
      * @param other  Type descriptor to cache the result for
      * @param result Result of the type check
      */
     void cacheTypeCheckResult(CacheableTypeDescriptor other, boolean result);
+
+    TypeCheckCacheKey getLookupKey();
 }

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/api/types/semtype/Context.java

@@ -24,7 +24,6 @@
 import io.ballerina.runtime.internal.types.semtype.MappingAtomicType;
 
 import java.util.ArrayList;
-import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -46,24 +45,9 @@ public final class Context {
     public final Map<Bdd, BddMemo> listMemo = new WeakHashMap<>();
     public final Map<Bdd, BddMemo> mappingMemo = new WeakHashMap<>();
     public final Map<Bdd, BddMemo> functionMemo = new WeakHashMap<>();
-    private static final int MAX_CACHE_SIZE = 100;
-    private final Map<CacheableTypeDescriptor, TypeCheckCache<CacheableTypeDescriptor>> typeCheckCacheMemo;
 
     private Context(Env env) {
         this.env = env;
-        this.typeCheckCacheMemo = createTypeCheckCacheMemo();
-    }
-
-    private static Map<CacheableTypeDescriptor, TypeCheckCache<CacheableTypeDescriptor>> createTypeCheckCacheMemo() {
-        // This is fine since this map is not going to get leaked out of the context and
-        // context is unique to a thread. So there will be no concurrent modifications
-        return new LinkedHashMap<>(MAX_CACHE_SIZE, 1f, true) {
-            @Override
-            protected boolean removeEldestEntry(
-                    Map.Entry<CacheableTypeDescriptor, TypeCheckCache<CacheableTypeDescriptor>> eldest) {
-                return size() > MAX_CACHE_SIZE;
-            }
-        };
     }
 
     public static Context from(Env env) {
@@ -141,11 +125,4 @@ public FunctionAtomicType functionAtomicType(Atom atom) {
         return env.functionAtomType(atom);
     }
 
-    public TypeCheckCache<CacheableTypeDescriptor> getTypeCheckCache(CacheableTypeDescriptor typeDescriptor) {
-        return typeCheckCacheMemo.computeIfAbsent(typeDescriptor, TypeCheckCache::new);
-    }
-
-    enum Phase {
-        INIT, TYPE_RESOLUTION, TYPE_CHECKING
-    }
 }

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/api/types/semtype/TypeCheckCache.java

@@ -1,41 +1,44 @@
 package io.ballerina.runtime.api.types.semtype;
 
-import io.ballerina.runtime.api.types.Type;
+import io.ballerina.runtime.internal.types.semtype.UniqueLookupKey;
 
 import java.util.Map;
 import java.util.Optional;
 import java.util.WeakHashMap;
 
 /**
- * Generalized implementation of type check result cache. It is okay to access
- * this from multiple threads but makes no
- * guarantee about the consistency of the cache under parallel access. Given
- * result don't change due to race conditions
+ * Generalized implementation of type check result cache. It is okay to access this from multiple threads but makes no
+ * guarantee about the consistency of the cache under parallel access. Given result don't change due to race conditions
  * this should eventually become consistent.
  *
- * @param <T> Type of the type descriptor which owns this cache
  * @since 2201.11.0
  */
-public class TypeCheckCache<T extends Type> {
+public class TypeCheckCache {
 
     // Not synchronizing this should be fine since race conditions don't lead to inconsistent results. (i.e. results
     // of doing multiple type checks are agnostic to the order of execution). Data races shouldn't lead to tearing in
     // 64-bit JVMs.
-    private final Map<T, Boolean> cachedResults = new WeakHashMap<>();
-    private final T owner;
+    private final Map<TypeCheckCacheKey, Boolean> cachedResults = new WeakHashMap<>();
+    private final TypeCheckCacheKey ownerKey;
 
-    public TypeCheckCache(T owner) {
-        this.owner = owner;
+    public TypeCheckCache(CacheableTypeDescriptor owner) {
+        this.ownerKey = owner.getLookupKey();
     }
 
-    public Optional<Boolean> cachedTypeCheckResult(T other) {
-        if (other.equals(owner)) {
+    public Optional<Boolean> cachedTypeCheckResult(CacheableTypeDescriptor other) {
+        TypeCheckCacheKey otherKey = other.getLookupKey();
+        if (otherKey.equals(ownerKey)) {
             return Optional.of(true);
         }
-        return Optional.ofNullable(cachedResults.get(other));
+        return Optional.ofNullable(cachedResults.get(otherKey));
     }
 
-    public void cacheTypeCheckResult(T other, boolean result) {
-        cachedResults.put(other, result);
+    public void cacheTypeCheckResult(CacheableTypeDescriptor other, boolean result) {
+        TypeCheckCacheKey lookupKey = other.getLookupKey();
+        // FIXME: this is just to prevent cache becoming too large. Revisit this after structured keys
+        if (lookupKey instanceof UniqueLookupKey) {
+            return;
+        }
+        cachedResults.put(lookupKey, result);
     }
 }

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/api/types/semtype/TypeCheckCacheFactory.java

@@ -0,0 +1,45 @@
+package io.ballerina.runtime.api.types.semtype;
+
+import io.ballerina.runtime.internal.types.semtype.UniqueLookupKey;
+
+import java.util.Map;
+import java.util.WeakHashMap;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+
+public final class TypeCheckCacheFactory {
+
+    private static final ReadWriteLock lock = new ReentrantReadWriteLock();
+
+    private static final Map<TypeCheckCacheKey, TypeCheckCache> cache = new WeakHashMap<>();
+
+    private TypeCheckCacheFactory() {
+    }
+
+    public static TypeCheckCache getTypeCheckCache(CacheableTypeDescriptor owner) {
+        var key = owner.getLookupKey();
+        if (key instanceof UniqueLookupKey) {
+            return new TypeCheckCache(owner);
+        }
+        lock.readLock().lock();
+        try {
+            if (cache.containsKey(key)) {
+                return cache.get(key);
+            }
+        } finally {
+            lock.readLock().unlock();
+        }
+        lock.writeLock().lock();
+        try {
+            if (cache.containsKey(key)) {
+                return cache.get(key);
+            }
+            var typeCheckCache = new TypeCheckCache(owner);
+            cache.put(key, typeCheckCache);
+            return typeCheckCache;
+        } finally {
+            lock.writeLock().unlock();
+        }
+    }
+
+}

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/api/types/semtype/TypeCheckCacheKey.java

@@ -0,0 +1,10 @@
+package io.ballerina.runtime.api.types.semtype;
+
+/**
+ * Marker type to represent the lookup key used in cache lookup operations involving {@CacheableTypeDescriptor}.
+ *
+ * @since 2201.12.0
+ */
+public interface TypeCheckCacheKey {
+
+}

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/internal/TypeChecker.java

@@ -639,13 +639,9 @@ private static boolean isSubTypeInner(Context cx, Type source, Type target) {
 
     private static boolean isSubTypeWithCache(Context cx, CacheableTypeDescriptor source,
                                               CacheableTypeDescriptor target) {
-        if (!source.shouldCache() || !target.shouldCache()) {
-            return isSubTypeInner(cx, source, target);
-        }
         Optional<Boolean> cachedResult = source.cachedTypeCheckResult(cx, target);
         logger.typeCheckCachedResult(cx, source, target, cachedResult);
         if (cachedResult.isPresent()) {
-            assert cachedResult.get() == isSubTypeInner(cx, source, target);
             return cachedResult.get();
         }
         boolean result = isSubTypeInner(cx, source, target);

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/internal/types/BAnyType.java

@@ -29,6 +29,7 @@
 import io.ballerina.runtime.api.types.semtype.ConcurrentLazySupplier;
 import io.ballerina.runtime.api.types.semtype.Core;
 import io.ballerina.runtime.api.types.semtype.SemType;
+import io.ballerina.runtime.internal.types.semtype.StructuredLookupKey;
 import io.ballerina.runtime.internal.values.RefValue;
 
 import java.util.Optional;
@@ -121,6 +122,11 @@ public void setImmutableType(IntersectionType immutableType) {
             this.immutableType = immutableType;
         }
 
+        @Override
+        public StructuredLookupKey getStructuredLookupKey() {
+            return null;
+        }
+
         @Override
         public Optional<IntersectionType> getIntersectionType() {
             return this.intersectionType == null ? Optional.empty() : Optional.of(this.intersectionType);

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/internal/types/BArrayType.java

@@ -26,15 +26,20 @@
 import io.ballerina.runtime.api.types.semtype.Env;
 import io.ballerina.runtime.api.types.semtype.SemType;
 import io.ballerina.runtime.api.types.semtype.ShapeAnalyzer;
+import io.ballerina.runtime.api.types.semtype.TypeCheckCacheKey;
 import io.ballerina.runtime.internal.TypeChecker;
 import io.ballerina.runtime.internal.types.semtype.CellAtomicType;
 import io.ballerina.runtime.internal.types.semtype.DefinitionContainer;
 import io.ballerina.runtime.internal.types.semtype.ListDefinition;
+import io.ballerina.runtime.internal.types.semtype.StructuredLookupKey;
 import io.ballerina.runtime.internal.values.AbstractArrayValue;
 import io.ballerina.runtime.internal.values.ArrayValue;
 import io.ballerina.runtime.internal.values.ArrayValueImpl;
 import io.ballerina.runtime.internal.values.ReadOnlyUtils;
 
+import java.lang.ref.Reference;
+import java.lang.ref.WeakReference;
+import java.util.Collections;
 import java.util.Optional;
 import java.util.Set;
 
@@ -69,6 +74,7 @@ public class BArrayType extends BType implements ArrayType, TypeWithShape {
     private int typeFlags;
     private final DefinitionContainer<ListDefinition> defn = new DefinitionContainer<>();
     private final DefinitionContainer<ListDefinition> acceptedTypeDefn = new DefinitionContainer<>();
+    private Reference<StructuredLookupKey> lookupKey;
     public BArrayType(Type elementType) {
         this(elementType, false);
     }
@@ -263,6 +269,27 @@ protected boolean isDependentlyTypedInner(Set<MayBeDependentType> visited) {
         return elementType instanceof MayBeDependentType eType && eType.isDependentlyTyped(visited);
     }
 
+    @Override
+    public StructuredLookupKey getStructuredLookupKey() {
+        if (lookupKey != null && lookupKey.get() != null) {
+            return lookupKey.get();
+        }
+        StructuredLookupKey structuredLookupKey = new StructuredLookupKey(StructuredLookupKey.Kind.ARRAY);
+        lookupKey = new WeakReference<>(structuredLookupKey);
+        if (size == -1) {
+            structuredLookupKey.setChildren(new TypeCheckCacheKey[]{
+                    new StructuredLookupKey(StructuredLookupKey.Kind.REST,
+                            new TypeCheckCacheKey[]{StructuredLookupKey.from(elementType)})});
+        } else {
+            assert size >= 0;
+            var elementKey = StructuredLookupKey.from(elementType);
+            // TODO: think of a better way to represent this since size could be very large
+            structuredLookupKey.setChildren(
+                    Collections.nCopies(size, elementKey).toArray(TypeCheckCacheKey[]::new));
+        }
+        return structuredLookupKey;
+    }
+
     @Override
     public Optional<SemType> inherentTypeOf(Context cx, ShapeSupplier shapeSupplier, Object object) {
         if (!couldInherentTypeBeDifferent()) {

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/internal/types/BBooleanType.java

@@ -25,6 +25,7 @@
 import io.ballerina.runtime.api.types.semtype.Builder;
 import io.ballerina.runtime.api.types.semtype.ConcurrentLazySupplier;
 import io.ballerina.runtime.api.types.semtype.SemType;
+import io.ballerina.runtime.internal.types.semtype.StructuredLookupKey;
 
 import java.util.function.Supplier;
 
@@ -86,5 +87,10 @@ public int getTag() {
         public boolean isReadOnly() {
             return true;
         }
+
+        @Override
+        public StructuredLookupKey getStructuredLookupKey() {
+            return null;
+        }
     }
 }

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/internal/types/BByteType.java

@@ -25,6 +25,7 @@
 import io.ballerina.runtime.api.types.semtype.Builder;
 import io.ballerina.runtime.api.types.semtype.ConcurrentLazySupplier;
 import io.ballerina.runtime.api.types.semtype.SemType;
+import io.ballerina.runtime.internal.types.semtype.StructuredLookupKey;
 
 import java.util.function.Supplier;
 
@@ -91,5 +92,10 @@ public boolean isReadOnly() {
         public BType clone() {
             return super.clone();
         }
+
+        @Override
+        public StructuredLookupKey getStructuredLookupKey() {
+            return null;
+        }
     }
 }

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/internal/types/BDecimalType.java

@@ -25,6 +25,7 @@
 import io.ballerina.runtime.api.types.semtype.Builder;
 import io.ballerina.runtime.api.types.semtype.ConcurrentLazySupplier;
 import io.ballerina.runtime.api.types.semtype.SemType;
+import io.ballerina.runtime.internal.types.semtype.StructuredLookupKey;
 import io.ballerina.runtime.internal.values.DecimalValue;
 
 import java.math.BigDecimal;
@@ -93,5 +94,10 @@ public boolean isReadOnly() {
         public BType clone() {
             return super.clone();
         }
+
+        @Override
+        public StructuredLookupKey getStructuredLookupKey() {
+            return null;
+        }
     }
 }

bvm/ballerina-runtime/src/main/java/io/ballerina/runtime/internal/types/BErrorType.java

@@ -28,8 +28,11 @@
 import io.ballerina.runtime.api.types.semtype.Context;
 import io.ballerina.runtime.api.types.semtype.Core;
 import io.ballerina.runtime.api.types.semtype.SemType;
+import io.ballerina.runtime.api.types.semtype.TypeCheckCacheKey;
 import io.ballerina.runtime.api.values.BError;
 import io.ballerina.runtime.internal.types.semtype.ErrorUtils;
+import io.ballerina.runtime.internal.types.semtype.StructuredLookupKey;
+import io.ballerina.runtime.internal.types.semtype.UniqueLookupKey;
 import io.ballerina.runtime.internal.values.ErrorValue;
 import io.ballerina.runtime.internal.values.MapValueImpl;
 
@@ -47,6 +50,8 @@ public class BErrorType extends BAnnotatableType implements ErrorType, TypeWithS
     public BTypeIdSet typeIdSet;
     private IntersectionType intersectionType = null;
     private volatile DistinctIdSupplier distinctIdSupplier;
+    private final StructuredLookupKey lookupKey = new StructuredLookupKey(StructuredLookupKey.Kind.ERROR,
+            new TypeCheckCacheKey[]{new UniqueLookupKey()});
 
     public BErrorType(String typeName, Module pkg, Type detailType) {
         super(typeName, pkg, ErrorValue.class);
@@ -155,6 +160,12 @@ protected boolean isDependentlyTypedInner(Set<MayBeDependentType> visited) {
                 mayBeDependentType.isDependentlyTyped(visited);
     }
 
+    @Override
+    public StructuredLookupKey getStructuredLookupKey() {
+        // TODO: need to think how to handle distinct types
+        return lookupKey;
+    }
+
     private boolean isTopType() {
         return detailType == PredefinedTypes.TYPE_DETAIL;
     }