Merge pull request #4428 from unisonweb/topic/test2

codebase2/codebase-sqlite/U/Codebase/Sqlite/Operations.hs

@@ -70,6 +70,8 @@ module U.Codebase.Sqlite.Operations
     -- ** type index
     Q.addTypeToIndexForTerm,
     termsHavingType,
+    filterTermsByReferenceHavingType,
+    filterTermsByReferentHavingType,
 
     -- ** type mentions index
     Q.addTypeMentionsToIndexForTerm,
@@ -1048,6 +1050,24 @@ termsHavingType cTypeRef =
       set <- traverse s2cReferentId sIds
       pure (Set.fromList set)
 
+filterTermsByReferenceHavingType :: C.TypeReference -> [C.Reference.Id] -> Transaction [C.Reference.Id]
+filterTermsByReferenceHavingType cTypeRef cTermRefIds =
+  runMaybeT (c2hReference cTypeRef) >>= \case
+    Nothing -> pure []
+    Just sTypeRef -> do
+      sTermRefIds <- traverse c2sReferenceId cTermRefIds
+      matches <- Q.filterTermsByReferenceHavingType sTypeRef sTermRefIds
+      traverse s2cReferenceId matches
+
+filterTermsByReferentHavingType :: C.TypeReference -> [C.Referent.Id] -> Transaction [C.Referent.Id]
+filterTermsByReferentHavingType cTypeRef cTermRefIds =
+  runMaybeT (c2hReference cTypeRef) >>= \case
+    Nothing -> pure []
+    Just sTypeRef -> do
+      sTermRefIds <- traverse c2sReferentId cTermRefIds
+      matches <- Q.filterTermsByReferentHavingType sTypeRef sTermRefIds
+      traverse s2cReferentId matches
+
 typeReferenceForTerm :: S.Reference.Id -> Transaction S.ReferenceH
 typeReferenceForTerm = Q.getTypeReferenceForReferent . C.Referent.RefId
 

codebase2/codebase-sqlite/U/Codebase/Sqlite/Queries.hs

@@ -167,6 +167,8 @@ module U.Codebase.Sqlite.Queries
     getReferentsByType,
     getTypeReferenceForReferent,
     getTypeReferencesForComponent,
+    filterTermsByReferenceHavingType,
+    filterTermsByReferentHavingType,
 
     -- ** type mentions index
     addToTypeMentionsIndex,
@@ -1460,6 +1462,76 @@ getTypeReferencesForComponent oId =
         WHERE term_referent_object_id = :oId
       |]
 
+filterTermsByReferentHavingType :: S.ReferenceH -> [S.Referent.Id] -> Transaction [S.Referent.Id]
+filterTermsByReferentHavingType typ terms = create *> for_ terms insert *> select <* drop
+  where
+    select = queryListRow [sql|
+      SELECT
+        q.term_referent_object_id,
+        q.term_referent_component_index,
+        q.term_referent_constructor_index
+      FROM filter_query q, find_type_index t
+      WHERE t.type_reference_builtin IS :typeBuiltin
+        AND t.type_reference_hash_id IS :typeHashId
+        AND t.type_reference_component_index IS :typeComponentIndex
+        AND t.term_referent_object_id = q.term_referent_object_id
+        AND t.term_referent_component_index = q.term_referent_component_index
+        AND t.term_referent_constructor_index IS q.term_referent_constructor_index
+    |]
+    insert r = execute [sql|
+      INSERT INTO filter_query (
+        term_referent_object_id,
+        term_referent_component_index,
+        term_referent_constructor_index
+      ) VALUES (@r, @, @)
+    |]
+    typeBuiltin :: Maybe TextId = Lens.preview C.Reference.t_ typ
+    typeHashId :: Maybe HashId = Lens.preview (C.Reference._ReferenceDerived . C.Reference.idH) typ
+    typeComponentIndex :: Maybe C.Reference.Pos = Lens.preview (C.Reference._ReferenceDerived . C.Reference.idPos) typ
+    create =  execute
+      [sql|
+        CREATE TEMPORARY TABLE filter_query (
+          term_referent_object_id INTEGER NOT NULL,
+          term_referent_component_index INTEGER NOT NULL,
+          term_referent_constructor_index INTEGER NULL
+        )
+      |]
+    drop =  execute [sql|DROP TABLE filter_query|]
+
+filterTermsByReferenceHavingType :: S.ReferenceH -> [S.Reference.Id] -> Transaction [S.Reference.Id]
+filterTermsByReferenceHavingType typ terms = create *> for_ terms insert *> select <* drop
+  where
+    select = queryListRow [sql|
+      SELECT
+        q.term_reference_object_id,
+        q.term_reference_component_index
+      FROM filter_query q, find_type_index t
+      WHERE t.type_reference_builtin IS :typeBuiltin
+        AND t.type_reference_hash_id IS :typeHashId
+        AND t.type_reference_component_index IS :typeComponentIndex
+        AND t.term_referent_object_id = q.term_reference_object_id
+        AND t.term_referent_component_index = q.term_reference_component_index
+        AND t.term_referent_constructor_index IS NULL
+    |]
+    insert r = execute [sql|
+      INSERT INTO filter_query (
+        term_reference_object_id,
+        term_reference_component_index
+      ) VALUES (@r, @)
+    |]
+    typeBuiltin :: Maybe TextId = Lens.preview C.Reference.t_ typ
+    typeHashId :: Maybe HashId = Lens.preview (C.Reference._ReferenceDerived . C.Reference.idH) typ
+    typeComponentIndex :: Maybe C.Reference.Pos = Lens.preview (C.Reference._ReferenceDerived . C.Reference.idPos) typ
+    create =  execute
+      [sql|
+        CREATE TEMPORARY TABLE filter_query (
+          term_reference_object_id INTEGER NOT NULL,
+          term_reference_component_index INTEGER NOT NULL
+        )
+      |]
+    drop =  execute [sql|DROP TABLE filter_query|]
+
+
 addToTypeMentionsIndex :: Reference' TextId HashId -> Referent.Id -> Transaction ()
 addToTypeMentionsIndex tp tm =
   execute

codebase2/core/U/Codebase/Reference.hs

@@ -10,7 +10,7 @@ module U.Codebase.Reference
     Reference' (..),
     TermReference',
     TypeReference',
-    ReferenceType(..),
+    ReferenceType (..),
     pattern Derived,
     Id,
     Id' (..),
@@ -20,6 +20,7 @@ module U.Codebase.Reference
     t_,
     h_,
     idH,
+    idPos,
     idToHash,
     idToShortHash,
     isBuiltin,
@@ -30,7 +31,7 @@ module U.Codebase.Reference
   )
 where
 
-import Control.Lens (Lens, Prism, Prism', Traversal, lens, preview, prism)
+import Control.Lens (Lens, Lens', Prism, Prism', Traversal, lens, preview, prism)
 import Data.Bifoldable (Bifoldable (..))
 import Data.Bitraversable (Bitraversable (..))
 import Data.Text qualified as Text
@@ -112,16 +113,19 @@ type Pos = Word64
 data Id' h = Id h Pos
   deriving stock (Eq, Ord, Show, Functor, Foldable, Traversable)
 
-t_ :: Traversal (Reference' t h) (Reference' t' h) t t'
-t_ f = \case
-  ReferenceBuiltin t -> ReferenceBuiltin <$> f t
-  ReferenceDerived id -> pure (ReferenceDerived id)
+t_ :: Prism (Reference' t h) (Reference' t' h) t t'
+t_ = prism ReferenceBuiltin \case
+  ReferenceBuiltin t -> Right t
+  ReferenceDerived id -> Left (ReferenceDerived id)
 
 h_ :: Traversal (Reference' t h) (Reference' t h') h h'
 h_ f = \case
   ReferenceBuiltin t -> pure (ReferenceBuiltin t)
   Derived h i -> Derived <$> f h <*> pure i
 
+idPos :: Lens' (Id' h) Pos
+idPos = lens (\(Id _h w) -> w) (\(Id h _w) w -> Id h w)
+
 idH :: Lens (Id' h) (Id' h') h h'
 idH = lens (\(Id h _w) -> h) (\(Id _h w) h -> Id h w)
 

parser-typechecker/src/Unison/Codebase.hs

@@ -20,6 +20,8 @@ module Unison.Codebase
 
     -- ** Search
     termsOfType,
+    filterTermsByReferenceIdHavingType,
+    filterTermsByReferentHavingType,
     termsMentioningType,
     SqliteCodebase.Operations.termReferencesByPrefix,
     termReferentsByPrefix,
@@ -155,7 +157,7 @@ import Unison.NameSegment qualified as NameSegment
 import Unison.Parser.Ann (Ann)
 import Unison.Parser.Ann qualified as Parser
 import Unison.Prelude
-import Unison.Reference (Reference)
+import Unison.Reference (Reference, TermReferenceId, TypeReference)
 import Unison.Reference qualified as Reference
 import Unison.Referent qualified as Referent
 import Unison.Runtime.IOSource qualified as IOSource
@@ -461,6 +463,28 @@ termsOfTypeByReference c r =
     . Set.map (fmap Reference.DerivedId)
     <$> termsOfTypeImpl c r
 
+filterTermsByReferentHavingType :: (Var v) => Codebase m v a -> Type v a -> Set Referent.Referent -> Sqlite.Transaction (Set Referent.Referent)
+filterTermsByReferentHavingType c ty = filterTermsByReferentHavingTypeByReference c $ Hashing.typeToReference ty
+
+filterTermsByReferenceIdHavingType :: (Var v) => Codebase m v a -> Type v a -> Set TermReferenceId -> Sqlite.Transaction (Set TermReferenceId)
+filterTermsByReferenceIdHavingType c ty = filterTermsByReferenceIdHavingTypeImpl c (Hashing.typeToReference ty)
+
+-- | Find the subset of `tms` which match the exact type `r` points to.
+filterTermsByReferentHavingTypeByReference :: Codebase m v a -> TypeReference -> Set Referent.Referent -> Sqlite.Transaction (Set Referent.Referent)
+filterTermsByReferentHavingTypeByReference c r tms = do
+  let (builtins, derived) = partitionEithers . map p $ Set.toList tms
+  let builtins' =
+        Set.intersection
+          (Set.fromList builtins)
+          (Rel.lookupDom r Builtin.builtinTermsByType)
+  derived' <- filterTermsByReferentIdHavingTypeImpl c r (Set.fromList derived)
+  pure $ builtins' <> Set.mapMonotonic Referent.fromId derived'
+  where
+    p :: Referent.Referent -> Either Referent.Referent Referent.Id
+    p r = case Referent.toId r of
+      Just rId -> Right rId
+      Nothing -> Left r
+
 -- | Get the set of terms-or-constructors mention the given type anywhere in their signature.
 termsMentioningType :: (Var v) => Codebase m v a -> Type v a -> Sqlite.Transaction (Set Referent.Referent)
 termsMentioningType c ty =

parser-typechecker/src/Unison/Codebase/SqliteCodebase.hs

@@ -71,7 +71,7 @@ import Unison.DataDeclaration (Decl)
 import Unison.Hash (Hash)
 import Unison.Parser.Ann (Ann)
 import Unison.Prelude
-import Unison.Reference (Reference)
+import Unison.Reference (Reference, TermReferenceId)
 import Unison.Reference qualified as Reference
 import Unison.Referent qualified as Referent
 import Unison.ShortHash (ShortHash)
@@ -352,6 +352,14 @@ sqliteCodebase debugName root localOrRemote lockOption migrationStrategy action
             termsOfTypeImpl =
               CodebaseOps.termsOfTypeImpl getDeclType
 
+            filterTermsByReferentIdHavingTypeImpl :: Reference -> Set Referent.Id -> Sqlite.Transaction (Set Referent.Id)
+            filterTermsByReferentIdHavingTypeImpl =
+              CodebaseOps.filterReferentsHavingTypeImpl getDeclType
+
+            filterTermsByReferenceIdHavingTypeImpl :: Reference -> Set TermReferenceId -> Sqlite.Transaction (Set TermReferenceId)
+            filterTermsByReferenceIdHavingTypeImpl =
+              CodebaseOps.filterReferencesHavingTypeImpl
+
             termsMentioningTypeImpl :: Reference -> Sqlite.Transaction (Set Referent.Id)
             termsMentioningTypeImpl =
               CodebaseOps.termsMentioningTypeImpl getDeclType
@@ -382,6 +390,8 @@ sqliteCodebase debugName root localOrRemote lockOption migrationStrategy action
                   getWatch,
                   termsOfTypeImpl,
                   termsMentioningTypeImpl,
+                  filterTermsByReferenceIdHavingTypeImpl,
+                  filterTermsByReferentIdHavingTypeImpl,
                   termReferentsByPrefix = referentsByPrefix,
                   withConnection = withConn,
                   withConnectionIO = withConnection debugName root

parser-typechecker/src/Unison/Codebase/SqliteCodebase/Conversions.hs

@@ -273,20 +273,16 @@ branchHash2to1 :: forall m. BranchHash -> V1.Branch.NamespaceHash m
 branchHash2to1 = V1.HashFor . unBranchHash
 
 reference2to1 :: V2.Reference -> V1.Reference
-reference2to1 = \case
-  V2.ReferenceBuiltin t -> V1.Reference.Builtin t
-  V2.ReferenceDerived i -> V1.Reference.DerivedId $ referenceid2to1 i
+reference2to1 = id
 
 reference1to2 :: V1.Reference -> V2.Reference
-reference1to2 = \case
-  V1.Reference.Builtin t -> V2.ReferenceBuiltin t
-  V1.Reference.DerivedId i -> V2.ReferenceDerived (referenceid1to2 i)
+reference1to2 = id
 
 referenceid1to2 :: V1.Reference.Id -> V2.Reference.Id
-referenceid1to2 (V1.Reference.Id h i) = V2.Reference.Id h i
+referenceid1to2 = id
 
 referenceid2to1 :: V2.Reference.Id -> V1.Reference.Id
-referenceid2to1 (V2.Reference.Id h i) = V1.Reference.Id h i
+referenceid2to1 = id
 
 rreferent2to1 :: (Applicative m) => Hash -> (V2.Reference -> m CT.ConstructorType) -> V2.ReferentH -> m V1.Referent
 rreferent2to1 h lookupCT = \case
@@ -314,6 +310,11 @@ referent1to2 = \case
   V1.Ref r -> V2.Ref $ reference1to2 r
   V1.Con (V1.ConstructorReference r i) _ct -> V2.Con (reference1to2 r) (fromIntegral i)
 
+referentid1to2 :: V1.Referent.Id -> V2.Referent.Id
+referentid1to2 = \case
+  V1.RefId r -> V2.RefId (referenceid1to2 r)
+  V1.ConId (V1.ConstructorReference r i) _ct -> V2.ConId (referenceid1to2 r) i
+
 referentid2to1 :: (Applicative m) => (V2.Reference -> m CT.ConstructorType) -> V2.Referent.Id -> m V1.Referent.Id
 referentid2to1 lookupCT = \case
   V2.RefId r -> pure $ V1.RefId (referenceid2to1 r)

parser-typechecker/src/Unison/Codebase/SqliteCodebase/Operations.hs

@@ -503,6 +503,23 @@ termsMentioningTypeImpl doGetDeclType r =
   Ops.termsMentioningType (Cv.reference1to2 r)
     >>= Set.traverse (Cv.referentid2to1 doGetDeclType)
 
+filterReferencesHavingTypeImpl :: Reference -> Set Reference.Id -> Transaction (Set Reference.Id)
+filterReferencesHavingTypeImpl typRef termRefs =
+  Ops.filterTermsByReferenceHavingType (Cv.reference1to2 typRef) (Cv.referenceid1to2 <$> toList termRefs)
+    <&> fmap Cv.referenceid2to1
+    <&> Set.fromList
+
+filterReferentsHavingTypeImpl ::
+  -- | A 'getDeclType'-like lookup, possibly backed by a cache.
+  (C.Reference.Reference -> Transaction CT.ConstructorType) ->
+  Reference ->
+  Set Referent.Id ->
+  Transaction (Set Referent.Id)
+filterReferentsHavingTypeImpl doGetDeclType typRef termRefs =
+  Ops.filterTermsByReferentHavingType (Cv.reference1to2 typRef) (Cv.referentid1to2 <$> toList termRefs)
+    >>= traverse (Cv.referentid2to1 doGetDeclType)
+      <&> Set.fromList
+
 -- | The number of base32 characters needed to distinguish any two references in the codebase.
 hashLength :: Transaction Int
 hashLength = pure 10

parser-typechecker/src/Unison/Codebase/Type.hs

@@ -27,7 +27,7 @@ import Unison.ConstructorType qualified as CT
 import Unison.DataDeclaration (Decl)
 import Unison.Hash (Hash)
 import Unison.Prelude
-import Unison.Reference (Reference)
+import Unison.Reference (Reference, TypeReference)
 import Unison.Reference qualified as Reference
 import Unison.Referent qualified as Referent
 import Unison.ShortHash (ShortHash)
@@ -98,6 +98,10 @@ data Codebase m v a = Codebase
     termsOfTypeImpl :: Reference -> Sqlite.Transaction (Set Referent.Id),
     -- | Get the set of user-defined terms-or-constructors mention the given type anywhere in their signature.
     termsMentioningTypeImpl :: Reference -> Sqlite.Transaction (Set Referent.Id),
+    -- | Return the subset of the given set that has the given type.
+    filterTermsByReferenceIdHavingTypeImpl :: TypeReference -> Set Reference.Id -> Sqlite.Transaction (Set Reference.Id),
+    -- | Return the subset of the given set that has the given type.
+    filterTermsByReferentIdHavingTypeImpl :: TypeReference -> Set Referent.Id -> Sqlite.Transaction (Set Referent.Id),
     -- | Get the set of user-defined terms-or-constructors whose hash matches the given prefix.
     termReferentsByPrefix :: ShortHash -> Sqlite.Transaction (Set Referent.Id),
     -- | Acquire a new connection to the same underlying database file this codebase object connects to.