introduce the concept of a singleton key for shared memory

Summary:
We had the pattern of only using one canonical key when we wanted to store sets in shared memory. We were previously relying on
ad-hoc int or string keys, but since we have 3 uses now, I think it makes sense to abstract this pattern.

Reviewed By: dkgi

Differential Revision: D14726545

fbshipit-source-id: d3f19be61adb8ddd1c26641b87480f89be783136

Author: sinancepel

ast/sharedMemory.ml

@@ -177,20 +177,20 @@ module HandleKeys = struct
     let description = "All handles"
   end
 
-  module HandleKeys = SharedMemory.WithCache (IntKey) (HandleKeysValue)
+  module HandleKeys = SharedMemory.WithCache (SharedMemory.SingletonKey) (HandleKeysValue)
 
   let get () =
-    HandleKeys.get 0
+    HandleKeys.get SharedMemory.SingletonKey.key
     |> Option.value ~default:File.Handle.Set.Tree.empty
 
   let clear () =
-    HandleKeys.remove_batch (HandleKeys.KeySet.singleton 0)
+    HandleKeys.remove_batch (HandleKeys.KeySet.singleton SharedMemory.SingletonKey.key)
 
   let add ~handles:new_keys =
     let handles = get () in
     clear ();
     let handles = File.Handle.Set.Tree.union handles new_keys in
-    HandleKeys.add 0 handles
+    HandleKeys.add SharedMemory.SingletonKey.key handles
 
   let remove ~handles:old_keys =
     let handles = get () in
@@ -200,7 +200,7 @@ module HandleKeys = struct
       |> File.Handle.Set.Tree.of_list
       |> File.Handle.Set.Tree.diff handles
     in
-    HandleKeys.add 0 handles
+    HandleKeys.add SharedMemory.SingletonKey.key handles
 
   let normalize () =
     let handles = get () in
@@ -209,16 +209,10 @@ module HandleKeys = struct
     |> File.Handle.Set.Tree.to_list
     |> List.sort ~compare:File.Handle.compare
     |> File.Handle.Set.Tree.of_list
-    |> HandleKeys.add 0
-
-  let hash_of_key =
-    HandleKeys.hash_of_key
-
-  let serialize_key =
-    HandleKeys.serialize_key
+    |> HandleKeys.add SharedMemory.SingletonKey.key
 
   let compute_hashes_to_keys () =
-    HandleKeys.compute_hashes_to_keys ~keys:[0]
+    HandleKeys.compute_hashes_to_keys ~keys:[SharedMemory.SingletonKey.key]
 end
 
 

ast/sharedMemory.mli

@@ -57,7 +57,7 @@ end
 
 module HandleKeys: sig
   module HandleKeysValue: Value.Type with type t = File.Handle.Set.Tree.t
-  module HandleKeys: module type of Memory.NoCache (IntKey) (HandleKeysValue)
+  module HandleKeys: module type of Memory.NoCache (Memory.SingletonKey) (HandleKeysValue)
 
   val get: unit -> File.Handle.Set.Tree.t
 
@@ -70,11 +70,6 @@ module HandleKeys: sig
 
   val normalize: unit -> unit
 
-  (* Exposed for testing. *)
-  val hash_of_key: int -> string
-  val serialize_key: int -> string
-
-
   val compute_hashes_to_keys: unit -> string String.Map.t
 end
 

ast/test/sharedMemoryTest.ml

@@ -66,7 +66,10 @@ let test_compute_hashes_to_keys _ =
        ~keys:[File.Handle.create "first.py"; File.Handle.create "second/__init__.py"]
     );
   assert_mapping_equal
-    [HandleKeys.hash_of_key 0, HandleKeys.serialize_key 0]
+    [
+      HandleKeys.HandleKeys.hash_of_key Memory.SingletonKey.key,
+      HandleKeys.HandleKeys.serialize_key Memory.SingletonKey.key;
+    ]
     (HandleKeys.compute_hashes_to_keys ());
   assert_mapping_equal
     [

command/test/serverTest.ml

@@ -1087,8 +1087,8 @@ let test_compute_hashes_to_keys context =
     @
     [
       to_binding
-        (OrderKeys.hash_of_key "Order")
-        (OrderKeys.serialize_key "Order");
+        (OrderKeys.hash_of_key SharedMemory.SingletonKey.key)
+        (OrderKeys.serialize_key SharedMemory.SingletonKey.key);
       to_binding
         (OrderIndices.hash_of_key (Type.Primitive "sixteen"))
         (OrderIndices.serialize_key (Type.Primitive "sixteen"));
@@ -1144,7 +1144,9 @@ let test_compute_hashes_to_keys context =
       to_binding
         (Coverage.SharedMemory.hash_of_key (File.Handle.create "sample.py"))
         (Coverage.SharedMemory.serialize_key (File.Handle.create "sample.py"));
-      to_binding (Protocols.hash_of_key 0) (Protocols.serialize_key 0);
+      to_binding
+        (Protocols.hash_of_key SharedMemory.SingletonKey.key)
+        (Protocols.serialize_key SharedMemory.SingletonKey.key);
     ]
     |> List.sort ~compare
   in

server/request.ml

@@ -457,10 +457,10 @@ let process_type_query_request ~state:({ State.environment; _ } as state) ~confi
           let map =
             Map.set
               String.Map.empty
-              ~key:(OrderKeys.hash_of_key "Order")
-              ~data:(OrderKeys.serialize_key "Order")
+              ~key:(OrderKeys.hash_of_key SharedMemory.SingletonKey.key)
+              ~data:(OrderKeys.serialize_key SharedMemory.SingletonKey.key)
           in
-          match OrderKeys.get "Order" with
+          match OrderKeys.get SharedMemory.SingletonKey.key with
           | Some indices ->
               let annotations = List.filter_map indices ~f:OrderAnnotations.get in
               extend_map
@@ -544,7 +544,10 @@ let process_type_query_request ~state:({ State.environment; _ } as state) ~confi
           extend_map map ~new_map:(Coverage.SharedMemory.compute_hashes_to_keys ~keys:handles)
         in
         (* Protocols. *)
-        let map = extend_map map ~new_map:(Protocols.compute_hashes_to_keys ~keys:[0]) in
+        let map =
+          extend_map
+            map
+            ~new_map:(Protocols.compute_hashes_to_keys ~keys:[SharedMemory.SingletonKey.key]) in
         map
         |> Map.to_alist
         |> List.sort ~compare:(fun (left, _) (right, _) -> String.compare left right)
@@ -710,7 +713,7 @@ let process_type_query_request ~state:({ State.environment; _ } as state) ~confi
                       Some {
                         TypeQuery.serialized_key;
                         kind = OrderKeyValue.description;
-                        actual_key = key;
+                        actual_key = Int.to_string key;
                         actual_value =
                           value
                           >>| List.to_string ~f:Int.to_string;

service/environment.ml

@@ -190,12 +190,12 @@ module SharedHandler: Analysis.Environment.Handler = struct
     ClassDefinitions.get
 
   let register_protocol protocol =
-    let protocols = Protocols.get 0 |> Option.value ~default:[] in
-    Protocols.add 0 (protocol :: protocols)
+    let protocols = Protocols.get SharedMemory.SingletonKey.key |> Option.value ~default:[] in
+    Protocols.add SharedMemory.SingletonKey.key (protocol :: protocols)
 
 
   let protocols () =
-    Protocols.get 0
+    Protocols.get SharedMemory.SingletonKey.key
     |> Option.value ~default:[]
 
   let register_module ~qualifier ~local_mode ~handle ~stub ~statements =
@@ -543,7 +543,7 @@ let populate_shared_memory
 
       add_table OrderIndices.write_through indices;
       add_table OrderAnnotations.write_through annotations;
-      OrderKeys.write_through "Order" (Hashtbl.keys annotations);
+      OrderKeys.write_through SharedMemory.SingletonKey.key (Hashtbl.keys annotations);
     in
     add_type_order order;
 
@@ -557,7 +557,7 @@ let populate_shared_memory
     add_table GlobalKeys.write_through (Hashtbl.map ~f:Hash_set.to_list global_keys);
     add_table DependentKeys.write_through (Hashtbl.map ~f:Hash_set.to_list dependent_keys);
 
-    Protocols.write_through 0 (Hash_set.to_list protocols);
+    Protocols.write_through SharedMemory.SingletonKey.key (Hash_set.to_list protocols);
     add_table
       (fun qualifier ast_module -> Ast.SharedMemory.Modules.add ~qualifier ~ast_module)
       modules;
@@ -584,13 +584,13 @@ let normalize_shared_memory () =
   (* Since we don't provide an API to the raw order keys in the type order handler,
      handle it inline here. *)
   begin
-    match OrderKeys.get "Order" with
+    match OrderKeys.get SharedMemory.SingletonKey.key with
     | None ->
         ()
     | Some keys ->
-        OrderKeys.remove_batch (OrderKeys.KeySet.singleton "Order");
+        OrderKeys.remove_batch (OrderKeys.KeySet.singleton SharedMemory.SingletonKey.key);
         List.sort ~compare:Int.compare keys
-        |> OrderKeys.add "Order";
+        |> OrderKeys.add SharedMemory.SingletonKey.key;
   end;
   Ast.SharedMemory.HandleKeys.normalize ();
   let handles = Ast.SharedMemory.HandleKeys.get () in

service/environmentSharedMemory.ml

@@ -148,7 +148,7 @@ module Globals = Memory.WithCache (Ast.SharedMemory.ReferenceKey) (GlobalValue)
 
 module Dependents = Memory.WithCache (Ast.SharedMemory.ReferenceKey) (DependentValue)
 
-module Protocols = Memory.WithCache (Ast.SharedMemory.IntKey) (ProtocolValue)
+module Protocols = Memory.WithCache (Memory.SingletonKey) (ProtocolValue)
 
 (** Keys *)
 module FunctionKeys = Memory.WithCache (Ast.SharedMemory.HandleKey) (FunctionKeyValue)
@@ -170,7 +170,7 @@ module OrderEdges = Memory.WithCache (Ast.SharedMemory.IntKey) (EdgeValue)
 
 module OrderBackedges = Memory.WithCache (Ast.SharedMemory.IntKey) (BackedgeValue)
 
-module OrderKeys = Memory.WithCache (StringKey) (OrderKeyValue)
+module OrderKeys = Memory.WithCache (Memory.SingletonKey) (OrderKeyValue)
 
 module StoredConfiguration = Memory.NoCache (StringKey) (ConfigurationValue)
 

service/memory.ml

@@ -231,3 +231,14 @@ let save_shared_memory ~path =
 
 let load_shared_memory ~path =
   SharedMem.load_table path
+
+
+module SingletonKey = struct
+  type t = int
+  let to_string = Core.Int.to_string
+  let compare = Core.Int.compare
+
+  type out = int
+  let from_string = Core.Int.of_string
+  let key = 0
+end

service/memory.mli

@@ -43,9 +43,9 @@ module NoCache (Key: KeyType) (Value: Value.Type): sig
 
   include SharedMemory.NoCache with
     type t = Value.t
-    and type key = Key.t
-    and module KeySet = Set.Make (Key)
-    and module KeyMap = MyMap.Make (Key)
+                                and type key = Key.t
+                                and module KeySet = Set.Make (Key)
+                                and module KeyMap = MyMap.Make (Key)
 end
 
 
@@ -58,9 +58,9 @@ module WithCache (Key: KeyType) (Value: Value.Type): sig
 
   include SharedMemory.WithCache with
     type t = Value.t
-    and type key = Key.t
-    and module KeySet = Set.Make (Key)
-    and module KeyMap = MyMap.Make (Key)
+                                  and type key = Key.t
+                                  and module KeySet = Set.Make (Key)
+                                  and module KeyMap = MyMap.Make (Key)
 end
 
 val get_heap_handle: Configuration.Analysis.t -> SharedMemory.handle
@@ -73,3 +73,14 @@ val save_shared_memory: path: string -> unit
 val load_shared_memory: path: string -> unit
 
 val unsafe_little_endian_representation: key: Digest.t -> Int64.t
+
+module SingletonKey: sig
+  type t
+  val to_string: t -> string
+  val compare: t -> t -> int
+
+  type out = int
+  val from_string: string -> int
+
+  val key: t
+end

service/test/environmentTest.ml

@@ -186,7 +186,7 @@ let test_normalize_order_keys _ =
   in
   Service.Environment.normalize_shared_memory ();
   assert_equal
-    (Service.EnvironmentSharedMemory.OrderKeys.get "Order")
+    (Service.EnvironmentSharedMemory.OrderKeys.get SharedMemory.SingletonKey.key)
     (Some indices)
 
 

service/typeOrder.ml

@@ -79,19 +79,19 @@ module Handler = struct
     set key data
 
   let length _ =
-    (OrderKeys.get "Order"
+    (OrderKeys.get SharedMemory.SingletonKey.key
      >>| List.length)
     |> Option.value ~default:0
 
   let add_key key =
-    match OrderKeys.get "Order" with
-    | None -> OrderKeys.add "Order" [key]
+    match OrderKeys.get SharedMemory.SingletonKey.key with
+    | None -> OrderKeys.add SharedMemory.SingletonKey.key [key]
     | Some keys ->
-        OrderKeys.remove_batch (OrderKeys.KeySet.singleton "Order");
-        OrderKeys.add "Order" (key :: keys)
+        OrderKeys.remove_batch (OrderKeys.KeySet.singleton SharedMemory.SingletonKey.key);
+        OrderKeys.add SharedMemory.SingletonKey.key (key :: keys)
 
   let keys () =
-    Option.value ~default:[] (OrderKeys.get "Order")
+    Option.value ~default:[] (OrderKeys.get SharedMemory.SingletonKey.key)
 
   let show () =
     let keys =