Overload the elements syntax to work for both arrays and sets (take 2)

Summary:
The biggest lesson from with migrating arrays to sets is having to change all uses of the elements syntax from X[..] to elements X. This also caused a lot of breakage for a variety of queries in production. The right way to fix this is to just overload the syntax for arrays to also work for sets. That's what this diff does.

I had a previous crack at this in D68559811 using a form of constraint machinery. This diff uses a more straightforward unification mechanism and is simpler.

Reviewed By: simonmar

Differential Revision: D68624547

fbshipit-source-id: aa42f879e54f94c434b7c07433c1be27fc6000d9

glean/db/Glean/Database/Schema/Types.hs

@@ -325,6 +325,7 @@ mkRtsType lookupType lookupPid = rtsType
     rtsType Schema.TyVar{} = error "rtsType: TyVar"
     rtsType Schema.HasTy{} = error "rtsType: HasTy"
     rtsType Schema.HasKey{} = error "rtsType: HasKey"
+    rtsType Schema.ElementsOf{} = error "rtsType: ElementsOf"
 
     fieldType :: Schema.FieldDef -> Maybe FieldDef
     fieldType (Schema.FieldDef name ty) = Schema.FieldDef name <$> rtsType ty

glean/db/Glean/Query/Expand.hs

@@ -100,6 +100,8 @@ instantiateWithFreshVariables query numVars = do
     TcElementsOfArray (instantiatePat base pat)
   instantiateTcTerm base (TcElementsOfSet pat) =
     TcElementsOfSet (instantiatePat base pat)
+  instantiateTcTerm base (TcElementsUnresolved ty pat) =
+    TcElementsUnresolved ty (instantiatePat base pat)
   instantiateTcTerm base (TcQueryGen query) =
     TcQueryGen (instantiateQuery base query)
   instantiateTcTerm base (TcAll query) =

glean/db/Glean/Query/JSON.hs

@@ -276,6 +276,7 @@ encode expanded Encoder{..} !d = enc
       TyVar{} -> error "JSON.encode: TyVar"
       HasTy{} -> error "JSON.encode: HasTy"
       HasKey{} -> error "JSON.encode: HasKey"
+      ElementsOf{} -> error "JSON.encode: ElementsOf"
 
     {-# INLINE enc_field #-}
     enc_field prev i name ty d = do

glean/db/Glean/Query/Prune.hs

@@ -153,6 +153,8 @@ prune hasFacts (QueryWithInfo q _ gen t) = do
           <$> prunePat x
         TcElementsOfSet x ->
           Ref . MatchExt . Typed ty . TcElementsOfSet <$> prunePat x
+        TcElementsUnresolved ty' x ->
+          Ref . MatchExt . Typed ty . TcElementsUnresolved ty' <$> prunePat x
         TcQueryGen q ->
           Ref . MatchExt . Typed ty . TcQueryGen <$> pruneTcQuery q
         -- we dont' want to handle negation here because if it tries to match
@@ -237,6 +239,7 @@ renumberVars gen ty q =
       TcFactGen ref <$> renamePat k <*> renamePat v <*> pure range
     TcElementsOfArray x -> TcElementsOfArray <$> renamePat x
     TcElementsOfSet x -> TcElementsOfSet <$> renamePat x
+    TcElementsUnresolved ty x -> TcElementsUnresolved ty <$> renamePat x
     TcQueryGen q -> TcQueryGen <$> renameQuery q
     TcAll query -> TcAll <$> renameQuery query
     TcNegation xs -> TcNegation <$> traverse renameStmt xs

glean/db/Glean/Query/Typecheck.hs

@@ -340,13 +340,11 @@ inferExpr ctx pat = case pat of
 
   ElementsOfArray _ e -> do
     (e', ty) <- inferExpr ContextExpr e
-    elemTy <- case ty of
-      ArrayTy elemTy -> return elemTy
-      _other -> do
-        elemTy <- freshTyVar
-        inPat pat $ unify ty (ArrayTy elemTy)
-        return elemTy
-    return (Ref (MatchExt (Typed elemTy (TcElementsOfArray e'))), elemTy)
+    elemTy <- freshTyVar
+    containerTy <- freshTyVarInt
+    unify ty (ElementsOf elemTy containerTy)
+    return (Ref (MatchExt (Typed elemTy (TcElementsUnresolved (TyVar containerTy) e')))
+           ,elemTy)
 
   Elements _ e -> do
     (e', ty) <- inferExpr ContextExpr e
@@ -1031,6 +1029,7 @@ tcQueryDeps q = Set.fromList $ map getRef (overQuery q)
       TcFactGen pref x y _ -> pref : overPat x <> overPat y
       TcElementsOfArray x -> overPat x
       TcElementsOfSet x -> overPat x
+      TcElementsUnresolved _ x -> overPat x
       TcQueryGen q -> overQuery q
       TcAll q -> overQuery q
       TcNegation stmts -> foldMap overStatement stmts
@@ -1087,6 +1086,7 @@ tcTermUsesNegation = \case
   TcFactGen _ x y _ -> tcPatUsesNegation x <|> tcPatUsesNegation y
   TcElementsOfArray x -> tcPatUsesNegation x
   TcElementsOfSet x -> tcPatUsesNegation x
+  TcElementsUnresolved _ p -> tcPatUsesNegation p
   TcQueryGen q -> tcQueryUsesNegation q
   TcAll query -> tcQueryUsesNegation query
   TcNegation _ -> Just PatternNegation

glean/db/Glean/Query/Typecheck/Types.hs

@@ -58,6 +58,7 @@ data TcTerm
   | TcFactGen PidRef TcPat TcPat SeekSection
   | TcElementsOfArray TcPat
   | TcElementsOfSet TcPat
+  | TcElementsUnresolved Type TcPat
   | TcQueryGen TcQuery
   | TcAll TcQuery
   | TcNegation [TcStatement]
@@ -103,6 +104,7 @@ instance Display TcTerm where
     isUnit _ = False
   display opts (TcElementsOfArray arr) = displayAtom opts arr <> "[..]"
   display opts (TcElementsOfSet set) = "elements" <+> parens (display opts set)
+  display opts (TcElementsUnresolved _ pat) = displayAtom opts pat <> "[..]"
   display opts (TcQueryGen q) = parens (display opts q)
   display opts (TcAll query)
     = "all" <+> "(" <> display opts query <> ")"

glean/db/Glean/Query/Typecheck/Unify.hs

@@ -147,6 +147,18 @@ unify (HasKey a x) (HasKey b y)
   extend y (HasKey a x)
 unify a@PredicateTy{} b@HasKey{} = unify b a
 
+unify (ElementsOf elemTy x) setTy@(SetTy ty) = do
+  extend x setTy
+  unify elemTy ty
+unify a@SetTy{} b@ElementsOf{} = unify b a
+unify (ElementsOf elemTy x) arrTy@(ArrayTy ty) = do
+  extend x arrTy
+  unify elemTy ty
+unify a@ArrayTy{} b@ElementsOf{} = unify b a
+unify (ElementsOf elemTyA avar) b@(ElementsOf elemTyB _) = do
+  unify elemTyA elemTyB
+  extend avar b
+
 unify a b = unifyError a b
 
 unifyError :: Type -> Type -> T a
@@ -250,6 +262,11 @@ apply_ unbound unboundHas t = do
           Nothing -> return t
           Just u -> go u
       SetTy t -> SetTy <$> go t
+      ElementsOf _ x -> do
+        m <- lookup x
+        case m of
+          Nothing -> return t
+          Just u -> go u
 
 zonkVars :: T ()
 zonkVars = do
@@ -336,6 +353,24 @@ zonkTcPat p = case p of
             , "pattern: " <> display opts p
             , "expected type: " <> display opts ty
             ]
+  Ref (MatchExt (Typed ty (TcElementsUnresolved containerTy p))) -> do
+    ty' <- zonkType ty
+    containerTy' <- zonkType containerTy
+    p' <- zonkTcPat p
+    case containerTy' of
+      SetTy _ ->
+        return (Ref (MatchExt (Typed ty' (TcElementsOfSet p'))))
+      ArrayTy _ ->
+        return  (Ref (MatchExt (Typed ty' (TcElementsOfArray p'))))
+      _other -> do
+        opts <- gets tcDisplayOpts
+        prettyError $
+          nest 4 $ vcat
+            [ "type error in pattern"
+            , "pattern: " <> display opts p
+            , "expected type: " <> display opts ty
+            ]
+
 
   Ref m -> Ref <$> zonkMatch m
 
@@ -388,6 +423,7 @@ zonkTcTerm t = case t of
       <$> (Typed <$> zonkType ty <*> zonkTcPat p)
       <*> pure f
   TcElementsOfSet p -> TcElementsOfSet <$> zonkTcPat p
+  TcElementsUnresolved{} -> error "zonkTcTerm: TcElementsUnresolved" -- handled in zonkTcPat
   TcPromote{} -> error "zonkTcTerm: TcPromote" -- handled in zonkTcPat
   TcDemote{} -> error "zonkTcTerm: TcPromote" -- handled in zonkTcPat
   TcStructPat{} -> error "zonkTcTerm: TcStructPat" -- handled in zonkTcPat

glean/hs/Glean/Angle/Types.hs

@@ -425,6 +425,9 @@ data Type_ pref tref
     -- HasKey K X
     --   A type variable X that is constrained to be a predicate
     --   type with key K
+  | ElementsOf (Type_ pref tref) {-# UNPACK #-}!Int
+    -- ElementsOf T X
+    --   A type variable that is either set T or [T]
   deriving (Eq, Show, Functor, Foldable, Generic)
 
 data RecordOrSum = Record | Sum
@@ -451,6 +454,7 @@ instance Bifunctor Type_ where
     TyVar x -> TyVar x
     HasTy m r x -> HasTy (bimap f g <$> m) r x
     HasKey ty x -> HasKey (bimap f g ty) x
+    ElementsOf ty x -> ElementsOf (bimap f g ty) x
 
 instance Bifoldable Type_ where
   bifoldr f g r = \case
@@ -469,6 +473,7 @@ instance Bifoldable Type_ where
     TyVar _ -> r
     HasTy m _ _ -> foldr (flip $ bifoldr f g) r m
     HasKey ty _ -> bifoldr f g r ty
+    ElementsOf ty _ -> bifoldr f g r ty
 
 {- Note [Types]
 
@@ -754,6 +759,8 @@ instance (Display pref, Display tref) => Display (Type_ pref tref) where
     doField (n, ty) = pretty n <> " : " <> display opts ty
   display opts (HasKey ty x) =
     "T" <> pretty x <> parens (display opts ty)
+  display opts (ElementsOf ty x) =
+    "T" <> pretty x <> brackets (display opts ty)
   displayAtom opts t = case t of
     MaybeTy{} -> parens $ display opts t
     EnumeratedTy{} -> parens $ display opts t

glean/hs/Glean/RTS/Types.hs

@@ -131,6 +131,7 @@ repType BooleanTy = repType lowerBool
 repType TyVar{} = error "repType: TyVar"
 repType HasTy{} = error "repType: HasTy"
 repType HasKey{} = error "repType: HasKey"
+repType ElementsOf{} = error "repType: ElementsOf"
 
 sumLike :: Type -> Maybe [Glean.RTS.Types.FieldDef]
 sumLike (SumTy fs) = Just fs

glean/schema/gen/Glean/Schema/Gen/Cpp.hs

@@ -295,6 +295,7 @@ reprTy here t = case t of
   TyVar{} -> error "reprTy: TyVar"
   HasTy{} -> error "reprTy: HasTy"
   HasKey{} -> error "reprTy: HasKey"
+  ElementsOf{} -> error "reprTy: ElementsOf"
 
 shareTypeDef :: NameSpaces -> ResolvedType -> CppGen Text
 shareTypeDef here t = do
@@ -349,6 +350,7 @@ valueTy here t = case t of
   TyVar{} -> error "valueTy: TyVar"
   HasTy{} -> error "valueTy: HasTy"
   HasKey{} -> error "valueTy: HasKey"
+  ElementsOf{} -> error "valueTy: ElementsOf"
 
 
 -- ----------------------------------------------------------------------------

glean/schema/gen/Glean/Schema/Gen/HackJson.hs

@@ -321,6 +321,7 @@ angleTypeInnerReprFor BooleanTy = return BooleanTBool
 angleTypeInnerReprFor TyVar{} = error "angleTypeInnerReprFor: TyVar"
 angleTypeInnerReprFor HasTy{} = error "angleTypeInnerReprFor: HasTy"
 angleTypeInnerReprFor HasKey{} = error "angleTypeInnerReprFor: HasKey"
+angleTypeInnerReprFor ElementsOf{} = error "angleTypeInnerReprFor: ElementsOf"
 
 defFile :: Either ResolvedPredicateDef ResolvedTypeDef -> FilePath
 defFile (Left p) = fileFor $ predicateDefName p
@@ -430,3 +431,4 @@ cyclesInDefs ctx defs = concatMap hasCycles sccs
   outEdgesT TyVar{} = error "outEdgesT: TyVar"
   outEdgesT HasTy{} = error "outEdgesT: HasTy"
   outEdgesT HasKey{} = error "outEdgesT: HasKey"
+  outEdgesT ElementsOf{} = error "outEdgesT: ElementsOf"

glean/schema/gen/Glean/Schema/Gen/Haskell.hs

@@ -261,6 +261,7 @@ haskellTy_ withId genSub here t = case t of
   TyVar{} -> error "haskellTy_: TyVar"
   HasTy{} -> error "haskellTy_: HasTy"
   HasKey{} -> error "haskellTy_: HasKey"
+  ElementsOf{} -> error "haskellTy_: ElementsOf"
 
 
 genPredicate :: ResolvedPredicateDef -> M [Text]

glean/schema/gen/Glean/Schema/Gen/OCaml.hs

@@ -267,6 +267,7 @@ genOCamlType ns namePolicy t = case t of
   TyVar{} -> error "genOCamlType: TyVar"
   HasTy{} -> error "genOCamlType: HasTy"
   HasKey{} -> error "genOCamlType: HasKey"
+  ElementsOf{} -> error "genOCamlType: ElementsOf"
   where
     genField fieldKind field = do
       ty <- genOCamlTypeFromField field ns namePolicy (fieldDefType field)
@@ -377,6 +378,7 @@ genOCamlToJson var ns namePolicy t = case t of
   TyVar{} -> error "genOCamlToJson: TyVar"
   HasTy{} -> error "genOCamlToJson: HasTy"
   HasKey{} -> error "genOCamlToJson: HasKey"
+  ElementsOf{} -> error "genOCamlToJson: ElementsOf"
 
 genOCamlToJsonFromField
   :: FieldDef_ PredicateRef TypeRef -> GenVars -> NameSpaces -> NamePolicy

glean/schema/gen/Glean/Schema/Gen/Thrift.hs

@@ -334,6 +334,7 @@ thriftTy here t = case t of
   TyVar{} -> error "thriftTy: TyVar"
   HasTy{} -> error "thriftTy: HasTy"
   HasKey{} -> error "thriftTy: HasKey"
+  ElementsOf{} -> error "thriftTy: ElementsOf"
 
 mkField :: [Text] -> Text -> Int -> Name -> Text -> Text
 mkField annots structOrUnion i p t =

glean/schema/gen/Glean/Schema/Gen/Utils.hs

@@ -136,6 +136,7 @@ orderDecls decls = map betterNotBeAnyCyclesIn sccs
   outEdgesT TyVar{} = error "outEdgesT: TyVar"
   outEdgesT HasTy{} = error "outEdgesT: HasTy"
   outEdgesT HasKey{} = error "outEdgesT: HasKey"
+  outEdgesT ElementsOf{} = error "outEdgesT: ElementsOf"
 
 {- Note [predicate type references]
 
@@ -473,3 +474,4 @@ addNamespaceDependencies nss =
   outEdgesT TyVar{} = error "outEdgesT: TyVar"
   outEdgesT HasTy{} = error "outEdgesT: HasTy"
   outEdgesT HasKey{} = error "outEdgesT: HasKey"
+  outEdgesT ElementsOf{} = error "outEdgesT: ELementsOf"

glean/shell/tests/integration/test-shell.t

@@ -66,6 +66,11 @@
 
 
 
+
+
+
+
+
   $ query "example.Class _"
   [>] example.Class _ (re)
   { "id": [0-9]+, "key": { "name": "Fish", "line": 30 } } (re)
@@ -77,6 +82,7 @@
 
 
 
+
   $ query "{ wrong = what } : string"
   [>] { wrong = what } : string (re)
   type error:
@@ -90,6 +96,7 @@
 
 
 
+
   $ query "_"
   [>] _ (re)
   query has ambiguous type
@@ -101,6 +108,7 @@
 
 
 
+
   $ query "A -> B"
   [>] A -> B (re)
   a key/value pattern (X -> Y) cannot be used here
@@ -111,6 +119,7 @@
 
 
 
+
   $ query "A -> B"
   [>] A -> B (re)
   a key/value pattern (X -> Y) cannot be used here
@@ -121,6 +130,7 @@
 
 
 
+
   $ query "A B"
   [>] A B (re)
   not in scope: A
@@ -131,6 +141,7 @@
 
 
 
+
   $ query "A B"
   [>] A B (re)
   not in scope: A
@@ -141,6 +152,7 @@
 
 
 
+
   $ query "B = 1; 1 = B"
   [>] B = 1; 1 = B (re)
   the last statement should be an expression: B = 1; 1 = B
@@ -151,19 +163,21 @@
 
 
 
+
   $ query "A = 1; B = A[..]"
   [>] A = 1; B = A[[]..[]] (re)
   type error:
     nat
   does not match:
-    [T1]
+    T2[T1]
 
   1 |  A = 1; B = A[..]
            ^
   [1]
 
 
 
+
   $ query "A = \"a\"; B = A : nat"
   [>] A = \"a\"; B = A : nat (re)
   type error:
@@ -177,6 +191,7 @@
 
 
 
+
   $ query "{ w = A } : { n : nat | s : nat }"
   [>] { w = A } : { n : nat | s : nat } (re)
   unknown alt: w
@@ -189,6 +204,7 @@
 
 
 
+
   $ query "{} : { n : nat | s : nat }"
   [>] {} : { n : nat | s : nat } (re)
   matching on a union type should have the form { field = pattern }
@@ -201,6 +217,7 @@
 
 
 
+
   $ query "A; A"
   [>] A; A (re)
   variable A has unknown type
@@ -219,6 +236,7 @@
 
 
 
+
 Recursive expansion is on by default
   $ "$GLEAN" --service "::1:$PORT" --minloglevel 10 shell --db "$DB" ":limit 1" "example.Parent _" | head -n -4
   [>] :limit 1 (re)