feat(shell/server): add search command

src/frontends/lean/completion.cpp

@@ -165,6 +165,59 @@ std::vector<json> get_decl_completions(std::string const & pattern, environment
     return completions;
 }
 
+void search_decls(std::string const & pattern, std::vector<pair<std::string, environment>> const & envs, options const & opts,
+                  std::vector<json> & completions) {
+    unsigned max_results = get_auto_completion_max_results(opts);
+    unsigned max_errors = get_fuzzy_match_max_errors(pattern.size());
+    name_map<pair<std::string, environment>> name2env;
+    std::vector<pair<name, name>> exact_matches;
+    std::vector<pair<std::string, name>> selected;
+    bitap_fuzzy_search matcher(pattern, max_errors);
+
+    for (auto & env_file : envs) {
+        auto & env = env_file.second;
+        env.for_each_declaration([&](declaration const & d) {
+            if (name2env.find(d.get_name())) return;
+            name2env.insert(d.get_name(), env_file);
+            if (is_projection(env, d.get_name())) {
+                auto s_name = d.get_name().get_prefix();
+            }
+            if (is_internal_name(d.get_name())) {
+                return;
+            }
+            if (auto it = exact_prefix_match(env, pattern, d)) {
+                exact_matches.emplace_back(*it, d.get_name());
+            } else {
+                std::string text = d.get_name().to_string();
+                if (matcher.match(text))
+                    selected.emplace_back(text, d.get_name());
+            }
+        });
+    }
+    unsigned num_results = 0;
+    if (!exact_matches.empty()) {
+        std::sort(exact_matches.begin(), exact_matches.end(),
+                  [](pair<name, name> const & p1, pair<name, name> const & p2) {
+                      return p1.first.size() < p2.first.size();
+                  });
+        for (pair<name, name> const & p : exact_matches) {
+            auto j = serialize_decl(p.first, p.second, name2env.find(p.second)->second, opts);
+            if (!j["source"].count("file"))
+                j["source"]["file"] = name2env.find(p.second)->first;
+            completions.push_back(j);
+            num_results++;
+            if (num_results >= max_results)
+                break;
+        }
+    }
+    filter_completions<name>(pattern, selected, completions, max_results - num_results, [&](name const & n) {
+        auto j = serialize_decl(n, name2env.find(n)->second, opts);
+        if (!j["source"].count("file"))
+            j["source"]["file"] = name2env.find(n)->first;
+        return j;
+    });
+}
+
 std::vector<json> get_field_completions(name const & s, std::string const & pattern, environment const & env, options const & opts) {
     std::vector<json> completions;
 

src/frontends/lean/completion.h

@@ -24,6 +24,9 @@ std::vector<json> get_interactive_tactic_completions(std::string const & pattern
 std::vector<json> get_attribute_completions(std::string const & pattern, environment const & env, options const & opts);
 std::vector<json> get_namespace_completions(std::string const & pattern, environment const & env, options const & opts);
 
+void search_decls(std::string const & pattern, std::vector<pair<std::string, environment>> const & envs,
+                  options const & opts, std::vector<json> & completions);
+
 void initialize_completion();
 void finalize_completion();
 }

src/library/module_mgr.cpp

@@ -19,6 +19,28 @@ Author: Gabriel Ebner
 
 namespace lean {
 
+environment module_info::get_latest_env() const {
+    if (m_snapshots) {
+        auto snap = *m_snapshots;
+        while (snap.m_next) {
+            if (auto next = peek(snap.m_next)) {
+                snap = *next;
+            } else {
+                break;
+            }
+        }
+        if (auto parser_snap = snap.m_snapshot_at_end) {
+            return parser_snap->m_env;
+        }
+    }
+    if (auto res = peek(m_result)) {
+        if (auto env = peek(res->m_loaded_module->m_env)) {
+            return *env;
+        }
+    }
+    return environment();
+}
+
 void module_mgr::mark_out_of_date(module_id const & id) {
     for (auto & mod : m_modules) {
         if (!mod.second || mod.second->m_out_of_date) continue;
@@ -371,6 +393,19 @@ std::vector<module_name> module_mgr::get_direct_imports(module_id const & id, st
     return imports;
 }
 
+std::vector<std::shared_ptr<module_info const>> module_mgr::get_all_modules() {
+    unique_lock<mutex> lock(m_mutex);
+
+    std::vector<std::shared_ptr<module_info const>> mods;
+    for (auto & mod : m_modules) {
+        if (mod.second) {
+            mods.push_back(mod.second);
+        }
+    }
+
+    return mods;
+}
+
 void module_mgr::cancel_all() {
     for (auto & m : m_modules) {
         if (auto mod = m.second) {

src/library/module_mgr.h

@@ -58,6 +58,8 @@ struct module_info {
     environment const & get_produced_env() const {
         return get(get(m_result).m_loaded_module->m_env);
     }
+
+    environment get_latest_env() const;
 };
 
 class module_vfs {
@@ -107,6 +109,8 @@ public:
 
     std::shared_ptr<module_info const> get_module(module_id const &);
 
+    std::vector<std::shared_ptr<module_info const>> get_all_modules();
+
     void cancel_all();
 
     void set_use_snapshots(bool use_snapshots) { m_use_snapshots = use_snapshots; }

src/shell/server.cpp

@@ -22,6 +22,7 @@ Authors: Gabriel Ebner, Leonardo de Moura, Sebastian Ullrich
 #include "frontends/lean/parser.h"
 #include "frontends/lean/info_manager.h"
 #include "frontends/lean/interactive.h"
+#include "frontends/lean/completion.h"
 #include "shell/server.h"
 
 namespace lean {
@@ -427,6 +428,8 @@ void server::handle_request(server::cmd_req const & req) {
         handle_async_response(req, handle_complete(req));
     } else if (command == "info") {
         handle_async_response(req, handle_info(req));
+    } else if (command == "search") {
+        send_msg(handle_search(req));
     } else if (command == "roi") {
         send_msg(handle_roi(req));
     } else if (command == "sleep") {
@@ -602,6 +605,23 @@ task<server::cmd_res> server::handle_info(server::cmd_req const & req) {
       .set_cancellation_token(m_bg_task_ctok).build();
 }
 
+server::cmd_res server::handle_search(server::cmd_req const & req) {
+    std::string query = req.m_payload.at("query");
+
+    std::vector<pair<std::string, environment>> envs_to_search;
+    for (auto & mod : m_mod_mgr->get_all_modules()) {
+        envs_to_search.emplace_back(mod->m_mod, mod->get_latest_env());
+    }
+
+    std::vector<json> results;
+    search_decls(query, envs_to_search, m_ios.get_options(), results);
+
+    json j;
+    j["results"] = results;
+
+    return cmd_res(req.m_seq_num, j);
+}
+
 std::tuple<std::string, module_src, time_t> server::load_module(module_id const & id, bool can_use_olean) {
     if (m_open_files.count(id)) {
         auto & ef = m_open_files[id];

src/shell/server.h

@@ -97,6 +97,7 @@ class server : public module_vfs {
     cmd_res handle_sync(cmd_req const & req);
     task<cmd_res> handle_complete(cmd_req const & req);
     task<cmd_res> handle_info(cmd_req const & req);
+    cmd_res handle_search(cmd_req const & req);
     cmd_res handle_roi(cmd_req const & req);
 
     json autocomplete(std::shared_ptr<module_info const> const & mod_info, bool skip_completions, pos_info const & pos);