From b27a3410bb91fd9b674de5dfe8fe10c2aa5fd7ef Mon Sep 17 00:00:00 2001
From: Kate <kit.ty.kate@disroot.org>
Date: Tue, 15 Aug 2023 20:00:23 +0100
Subject: [PATCH] Speedup OpamTreeCommand.run

Co-authored-by: R. Boujbel <rjbou@ocamlpro.com>
---
 src/client/opamTreeCommand.ml | 22 ++++++++++++++--------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/src/client/opamTreeCommand.ml b/src/client/opamTreeCommand.ml
index f3504948cb6..5cc853d6584 100644
--- a/src/client/opamTreeCommand.ml
+++ b/src/client/opamTreeCommand.ml
@@ -429,14 +429,16 @@ let raw_state tog st install =
   simulate_new_state tog st universe install names
 
 let run st tog ?no_constraint ?(no_switch=false) mode filter atoms =
+  let open OpamPackage.Set.Op in
   let select, missing =
     List.fold_left (fun (select, missing) atom ->
-        if OpamPackage.Set.disjoint
-            (OpamSwitchState.packages_of_atoms st [atom])
-            st.installed
-        then (select, atom :: missing)
-        else (fst atom :: select, missing)
-      ) ([], []) atoms
+        let installed =
+          OpamPackage.Set.filter (OpamFormula.check atom) st.installed
+        in
+        if OpamPackage.Set.is_empty installed then
+          (select, atom :: missing)
+        else (installed ++ select, missing))
+      (OpamPackage.Set.empty, []) atoms
   in
   let st, universe =
     let universe = get_universe tog st in
@@ -455,7 +457,7 @@ let run st tog ?no_constraint ?(no_switch=false) mode filter atoms =
           (match missing with | [_] -> "" | _ -> "s")
           (OpamStd.Format.pretty_list
              (List.map (fun (name, _) -> OpamPackage.Name.to_string name) missing));
-      if select = [] && atoms <> [] then
+      if OpamPackage.Set.is_empty select && atoms <> [] then
         OpamConsole.error_and_exit `Not_found "No package to display"
       else
         st, universe
@@ -463,7 +465,11 @@ let run st tog ?no_constraint ?(no_switch=false) mode filter atoms =
   if OpamPackage.Set.is_empty st.installed then
     OpamConsole.error_and_exit `Not_found "No package is installed"
   else
-  let forest = build st universe tog mode filter (select @ List.map fst missing) in
+  let simulated = OpamFormula.packages_of_atoms st.installed missing in
+  let forest =
+    build st universe tog mode filter
+      (OpamPackage.Name.Set.to_list (OpamPackage.names_of_packages (select ++ simulated)))
+  in
   print ?no_constraint forest;
   if OpamClientConfig.(!r.json_out) <> None then
     (if not no_switch then