purge protocols per handle

Summary:
We had a bug where we were never removing elements from the master list of protocols. This meant that if you registered a class as a protocol and then removed the protocol superclass, Pyre would still think of it as a protocol. To prevent the issue, keep track of protocols as a set and remove protocols per handle in the `purge` function.

Moved the DependencyHandler since we're relying on it to register protocols now.

Caught by the consistency checker.

Reviewed By: dkgi

Differential Revision: D14736415

fbshipit-source-id: 7beb9cf6869515b43c9f5c435bb854cd85da84f5

Author: sinancepel

analysis/dependencies.ml

@@ -15,6 +15,7 @@ type index = {
   alias_keys: (Type.t Hash_set.t) File.Handle.Table.t;
   global_keys: (Reference.t Hash_set.t) File.Handle.Table.t;
   dependent_keys: (Reference.t Hash_set.t) File.Handle.Table.t;
+  protocol_keys: (Type.t Hash_set.t) File.Handle.Table.t;
 }
 
 
@@ -30,6 +31,7 @@ module type Handler = sig
   val add_alias_key: handle: File.Handle.t -> Type.t -> unit
   val add_global_key: handle: File.Handle.t -> Reference.t -> unit
   val add_dependent_key: handle: File.Handle.t -> Reference.t -> unit
+  val add_protocol_key: handle: File.Handle.t -> Type.t -> unit
 
   val add_dependent: handle: File.Handle.t -> Reference.t -> unit
   val dependents: Reference.t -> File.Handle.Set.Tree.t option
@@ -39,6 +41,7 @@ module type Handler = sig
   val get_alias_keys: handle: File.Handle.t -> Type.t list
   val get_global_keys: handle: File.Handle.t -> Reference.t list
   val get_dependent_keys: handle: File.Handle.t -> Reference.t list
+  val get_protocol_keys: handle: File.Handle.t -> Type.t list
 
   val clear_keys_batch: File.Handle.t list -> unit
 
@@ -47,7 +50,7 @@ end
 
 
 let handler {
-    index = { function_keys; class_keys; alias_keys; global_keys; dependent_keys };
+    index = { function_keys; class_keys; alias_keys; global_keys; dependent_keys; protocol_keys };
     dependents;
   } =
   (module struct
@@ -105,6 +108,15 @@ let handler {
       | Some hash_set ->
           Hash_set.add hash_set dependent
 
+    let add_protocol_key ~handle protocol =
+      match Hashtbl.find protocol_keys handle with
+      | None ->
+          Hashtbl.set
+            protocol_keys
+            ~key:handle
+            ~data:(Type.Hash_set.of_list [protocol])
+      | Some hash_set ->
+          Hash_set.add hash_set protocol
 
     let add_dependent ~handle dependent =
       add_dependent_key ~handle dependent;
@@ -154,12 +166,19 @@ let handler {
       |> Option.value ~default:[]
 
 
+    let get_protocol_keys ~handle =
+      Hashtbl.find protocol_keys handle
+      >>| Hash_set.to_list
+      |> Option.value ~default:[]
+
+
     let clear_keys_batch handles =
       List.iter ~f:(Hashtbl.remove function_keys) handles;
       List.iter ~f:(Hashtbl.remove class_keys) handles;
       List.iter ~f:(Hashtbl.remove alias_keys) handles;
       List.iter ~f:(Hashtbl.remove global_keys) handles;
-      List.iter ~f:(Hashtbl.remove dependent_keys) handles
+      List.iter ~f:(Hashtbl.remove dependent_keys) handles;
+      List.iter ~f:(Hashtbl.remove protocol_keys) handles
 
 
     let normalize handles =
@@ -186,13 +205,14 @@ let create () =
     alias_keys = File.Handle.Table.create ();
     global_keys = File.Handle.Table.create ();
     dependent_keys = File.Handle.Table.create ();
+    protocol_keys = File.Handle.Table.create ();
   }
   in
   { index = index; dependents = Reference.Table.create () }
 
 
 let copy {
-    index = { function_keys; class_keys; alias_keys; global_keys; dependent_keys };
+    index = { function_keys; class_keys; alias_keys; global_keys; dependent_keys; protocol_keys };
     dependents } =
   {
     index = {
@@ -201,6 +221,7 @@ let copy {
       alias_keys = Hashtbl.copy alias_keys;
       global_keys = Hashtbl.copy global_keys;
       dependent_keys = Hashtbl.copy dependent_keys;
+      protocol_keys = Hashtbl.copy protocol_keys;
     };
     dependents = Hashtbl.copy dependents;
   }

analysis/dependencies.mli

@@ -14,6 +14,7 @@ type index = {
   alias_keys: (Type.t Hash_set.t) File.Handle.Table.t;
   global_keys: (Reference.t Hash_set.t) File.Handle.Table.t;
   dependent_keys: (Reference.t Hash_set.t) File.Handle.Table.t;
+  protocol_keys: (Type.t Hash_set.t) File.Handle.Table.t;
 }
 
 type t = {
@@ -27,6 +28,7 @@ module type Handler = sig
   val add_alias_key: handle: File.Handle.t -> Type.t -> unit
   val add_global_key: handle: File.Handle.t -> Reference.t -> unit
   val add_dependent_key: handle: File.Handle.t -> Reference.t -> unit
+  val add_protocol_key: handle: File.Handle.t -> Type.t -> unit
 
   val add_dependent: handle: File.Handle.t -> Reference.t -> unit
 
@@ -37,6 +39,7 @@ module type Handler = sig
   val get_alias_keys: handle: File.Handle.t -> Type.t list
   val get_global_keys: handle: File.Handle.t -> Reference.t list
   val get_dependent_keys: handle: File.Handle.t -> Reference.t list
+  val get_protocol_keys: handle: File.Handle.t -> Type.t list
 
   val clear_keys_batch: File.Handle.t list -> unit
 

analysis/environment.ml

@@ -35,7 +35,7 @@ module type Handler = sig
 
   val class_definition: Type.t -> Resolution.class_representation option
 
-  val register_protocol: Type.t -> unit
+  val register_protocol: handle: File.Handle.t -> Type.t -> unit
   val protocols: unit -> Type.t list
 
   val register_module
@@ -269,6 +269,8 @@ let handler
       |> purge_table_given_keys globals;
       List.concat_map ~f:(fun handle -> DependencyHandler.get_dependent_keys ~handle) handles
       |> purge_dependents;
+      List.concat_map ~f:(fun handle -> DependencyHandler.get_protocol_keys ~handle) handles
+      |> List.iter ~f:(Hash_set.remove protocols);
       DependencyHandler.clear_keys_batch handles;
       List.map handles ~f:(fun handle -> Source.qualifier ~handle)
       |> List.iter ~f:(Hashtbl.remove modules);
@@ -281,7 +283,8 @@ let handler
     let class_definition =
       Hashtbl.find class_definitions
 
-    let register_protocol protocol =
+    let register_protocol ~handle protocol =
+      DependencyHandler.add_protocol_key ~handle protocol;
       Hash_set.add protocols protocol
 
     let protocols () =
@@ -403,7 +406,9 @@ let register_class_definitions (module Handler: Handler) source =
             let annotated = Annotated.Class.create { Node.location; value = definition } in
 
             if Annotated.Class.is_protocol annotated then
-              Handler.register_protocol primitive;
+              begin
+                Handler.register_protocol ~handle primitive
+              end;
 
             Handler.set_class_definition
               ~primitive

analysis/environment.mli

@@ -36,7 +36,7 @@ module type Handler = sig
 
   val class_definition: Type.t -> Resolution.class_representation option
 
-  val register_protocol: Type.t -> unit
+  val register_protocol: handle: File.Handle.t -> Type.t -> unit
   val protocols: unit -> Type.t list
 
   val register_module

analysis/test/environmentTest.ml

@@ -1490,6 +1490,7 @@ let test_purge _ =
   let ((module Handler: Environment.Handler) as handler) = Environment.handler environment in
   let source = {|
       import a
+      class P(typing.Protocol): pass
       class baz.baz(): pass
       _T = typing.TypeVar("_T")
       x = 5
@@ -1511,11 +1512,15 @@ let test_purge _ =
   assert_equal
     (dependencies "a.py")
     (Some ["test.py"]);
+  assert_equal
+    (Handler.protocols ())
+    [Type.Primitive "P"];
 
   Handler.purge [File.Handle.create "test.py"];
 
   assert_is_none (Handler.class_definition (Type.Primitive "baz.baz"));
   assert_is_none (Handler.aliases (Type.Primitive "_T"));
+  assert_equal (Handler.protocols ()) [];
   assert_equal
     (dependencies "a.py")
     (Some [])

server/request.ml

@@ -626,6 +626,7 @@ let process_type_query_request ~state:({ State.environment; _ } as state) ~confi
                         actual_key = Int.to_string key;
                         actual_value =
                           value
+                          >>| Type.Set.Tree.to_list
                           >>| List.to_string ~f:Type.show;
                       }
                   | Ok (FunctionKeys.Decoded (key, value)) ->