Optimize mnesia cache querying

lib/nostrum/cache/guild_cache/mnesia.ex

@@ -291,7 +291,20 @@ if Code.ensure_loaded?(:mnesia) do
     @doc "Get a QLC handle for the guild cache."
     @spec query_handle :: :qlc.query_handle()
     def query_handle do
-      ms = [{{:_, :"$1", :"$2"}, [], [{{:"$1", :"$2"}}]}]
+      query_handle([])
+    end
+
+    @doc """
+    Retrieve a query handle for the table with optional match specification guards.
+
+    ## Match specification variables
+
+    - `$1`: The guild ID
+    - `$2` The guild object itself
+    """
+    @doc since: "0.10.0"
+    def query_handle(guards) do
+      ms = [{{:_, :"$1", :"$2"}, guards, [{{:"$1", :"$2"}}]}]
       :mnesia.table(@table_name, {:traverse, {:select, ms}})
     end
 

lib/nostrum/cache/member_cache/mnesia.ex

@@ -147,7 +147,23 @@ if Code.ensure_loaded?(:mnesia) do
     @doc "Get a QLC query handle for the member cache."
     @spec query_handle :: :qlc.query_handle()
     def query_handle do
-      :mnesia.table(@table_name)
+      query_handle([])
+    end
+
+    @doc """
+    Retrieve a query handle for the table with optional match specification guards.
+
+    ## Match specification variables
+
+    - `$1`: The guild ID the presence is stored on
+    - `$2`: The user associated with the presence
+    - `$3`: The member itself
+    """
+    @doc since: "0.10.0"
+    @spec query_handle([term()]) :: :qlc.query_handle()
+    def query_handle(guards) do
+      ms = [{{_name = :_, _key = :_, :"$1", :"$2", :"$3"}, guards, [{{{{:"$1", :"$2"}}, :"$3"}}]}]
+      :mnesia.table(@table_name, {:traverse, {:select, ms}})
     end
 
     @impl MemberCache

lib/nostrum/cache/presence_cache/mnesia.ex

@@ -100,8 +100,22 @@ if Code.ensure_loaded?(:mnesia) do
     @doc since: "0.8.0"
     @spec query_handle :: :qlc.query_handle()
     def query_handle do
-      # Note: :"$1" holds a pair here (`{guild_id, user_id}`).
-      ms = [{{:_, :"$1", :"$2"}, [], [{{:"$1", :"$2"}}]}]
+      query_handle([])
+    end
+
+    @doc """
+    Retrieve a query handle for the table with optional match specification guards.
+
+    ## Match specification variables
+
+    - `$1`: The guild ID the presence is stored on
+    - `$2`: The user associated with the presence
+    - `$3`: The user presence itself
+    """
+    @doc since: "0.10.0"
+    @spec query_handle([term()]) :: :qlc.query_handle()
+    def query_handle(guards) do
+      ms = [{{:_, {:"$1", :"$2"}, :"$3"}, guards, [{{{{:"$1", :"$2"}}, :"$3"}}]}]
       :mnesia.table(@table_name, {:traverse, {:select, ms}})
     end
 

src/nostrum_guild_cache_qlc.erl

@@ -10,5 +10,9 @@ all(Cache) ->
 
 
 get(RequestedGuildId, Cache) ->
-    qlc:q([Guild || {GuildId, Guild} <- Cache:query_handle(),
+    QH = case erlang:function_exported(Cache, query_handle, 1) of
+             true -> Cache:query_handle([{'==', '$1', {const, RequestedGuildId}}]);
+             false -> Cache:query_handle()
+         end,
+    qlc:q([Guild || {GuildId, Guild} <- QH,
                     GuildId =:= RequestedGuildId]).

src/nostrum_member_cache_qlc.erl

@@ -21,52 +21,53 @@
 %
 % Implement https://github.com/erlang/otp/issues/7268 to improve it.
 
--define(MNESIA_CACHE, 'Elixir.Nostrum.Cache.MemberCache.Mnesia').
-
-% These must be selected carefully so that QLC can plan using the indices properly.
--define(MNESIA_FORMAT, {_Tag, {GuildId, MemberId}, GuildId, MemberId, Member}).
--define(MNESIA_FORMAT_NOMEMBERID, {_Tag, {_GuildId, _MemberId}, GuildId, _, Member}).
-
 
 -spec by_user('Elixir.Nostrum.Struct.User':id(), module()) -> qlc:query_handle().
-by_user(UserId, ?MNESIA_CACHE) ->
-    qlc:q([{GuildId, Member} || ?MNESIA_FORMAT <- ?MNESIA_CACHE:query_handle(),
-                                MemberId =:= UserId]);
-
 by_user(UserId, Cache) ->
-    qlc:q([{GuildId, Member} || {{GuildId, MemberId}, Member} <- Cache:query_handle(),
+    QH = case erlang:function_exported(Cache, query_handle, 1) of
+             true -> Cache:query_handle([{'==', '$2', {const, UserId}}]);
+             false -> Cache:query_handle()
+         end,
+
+    qlc:q([{GuildId, Member} || {{GuildId, MemberId}, Member} <- QH,
                                 MemberId =:= UserId]).
 
 -spec by_guild('Elixir.Nostrum.Struct.Guild':id(), module()) -> qlc:query_handle().
-by_guild(RequestedGuildId, ?MNESIA_CACHE) ->
-    qlc:q([Member || ?MNESIA_FORMAT_NOMEMBERID <- ?MNESIA_CACHE:query_handle(),
-                     GuildId =:= RequestedGuildId]);
-
 by_guild(RequestedGuildId, Cache) ->
-    qlc:q([Member || {{GuildId, _MemberId}, Member} <- Cache:query_handle(),
+    QH = case erlang:function_exported(Cache, query_handle, 1) of
+             true -> Cache:query_handle([{'==', '$1', {const, RequestedGuildId}}]);
+             false -> Cache:query_handle()
+         end,
+
+    qlc:q([Member || {{GuildId, _MemberId}, Member} <- QH,
                      GuildId =:= RequestedGuildId]).
 
--spec lookup('Elixir.Nostrum.Struct.Guild':id(), 'Elixir.Nostrum.Struct.User':id(), module()) -> qlc:query_handle().
-lookup(RequestedGuildId, RequestedUserId, ?MNESIA_CACHE) ->
-    qlc:q([Member || ?MNESIA_FORMAT <- ?MNESIA_CACHE:query_handle(),
-                     GuildId =:= RequestedGuildId,
-                     MemberId =:= RequestedUserId]);
+-spec by_guild_with_user_id('Elixir.Nostrum.Struct.Guild':id(), module()) -> qlc:query_handle().
+by_guild_with_user_id(RequestedGuildId, Cache) ->
+    QH = case erlang:function_exported(Cache, query_handle, 1) of
+             true -> Cache:query_handle([{'==', '$1', {const, RequestedGuildId}}]);
+             false -> Cache:query_handle()
+         end,
+
+    qlc:q([{MemberId, Member} || {{GuildId, MemberId}, Member} <- QH,
+                                 GuildId =:= RequestedGuildId]).
 
+-spec lookup('Elixir.Nostrum.Struct.Guild':id(), 'Elixir.Nostrum.Struct.User':id(), module()) -> qlc:query_handle().
 lookup(RequestedGuildId, RequestedUserId, Cache) ->
-    qlc:q([Member || {{GuildId, MemberId}, Member} <- Cache:query_handle(),
+    QH = case erlang:function_exported(Cache, query_handle, 1) of
+             true ->
+                 Cache:query_handle([{'==', '$1', {const, RequestedGuildId}},
+                                     {'==', '$2', {const, RequestedUserId}}]);
+             false -> Cache:query_handle()
+         end,
+
+    qlc:q([Member || {{GuildId, MemberId}, Member} <- QH,
                      GuildId =:= RequestedGuildId,
                      MemberId =:= RequestedUserId]).
 
 
 -spec get_with_users('Elixir.Nostrum.Struct.Guild':id(), module(), module()) -> qlc:query_handle().
-get_with_users(RequestedGuildId, ?MNESIA_CACHE, UserCache) ->
-    qlc:q([{Member, User} || ?MNESIA_FORMAT <- ?MNESIA_CACHE:query_handle(),
-                          GuildId =:= RequestedGuildId,
-                          {UserId, User} <- UserCache:query_handle(),
-                          MemberId =:= UserId]);
-
 get_with_users(RequestedGuildId, MemberCache, UserCache) ->
-    qlc:q([{Member, User} || {{GuildId, MemberId}, Member} <- MemberCache:query_handle(),
-                          GuildId =:= RequestedGuildId,
+    qlc:q([{Member, User} || {MemberId, Member} <- by_guild_with_user_id(RequestedGuildId, MemberCache),
                           {UserId, User} <- UserCache:query_handle(),
                           MemberId =:= UserId]).

src/nostrum_presence_cache_qlc.erl

@@ -6,6 +6,14 @@
 % Optimized presence cache QLC queries.
 
 get(RequestedGuildId, RequestedUserId, Cache) ->
-    qlc:q([Presence || {{GuildId, UserId}, Presence} <- Cache:query_handle(),
+    QH = case erlang:function_exported(Cache, query_handle, 1) of
+             true ->
+                 Guards = [{'==', '$1', {const, RequestedGuildId}},
+                           {'==', '$2', {const, RequestedUserId}}],
+                 Cache:query_handle(Guards);
+             false ->
+                 Cache:query_handle()
+         end,
+    qlc:q([Presence || {{GuildId, UserId}, Presence} <- QH,
                                    GuildId =:= RequestedGuildId,
                                    UserId =:= RequestedUserId]).