Use a workspace level TokenMap and Engines (#7139)

## Description
This PR moves `TokenMap` and `Engines` out of the `Session` type up to
`ServerState` so a single instance of each can be shared across all
workspace members.

Up until now, each workspace member used a unique `Engines`. This meant
the standard library and any other dependancies were unable to utilise
the `QueryEngine` cache, leading to lots of duplicate entries. This was
leading to inflated initialization times when opening a new workspace
member and significantly increasing the RAM footprint of the language
server.

For example, opening each workspace member in the FUSD projects
workspace would see RAM usage exceeding ~30.0 gig. Now we no longer
exceed 2.0 gig of RAM.

Future PR's will move the `BuildPlanCache` up to the `ServerState` level
as well.

part of #7111 & #6226
closes #5645
closes #5856

## Checklist

- [x] I have linked to any relevant issues.
- [x] I have commented my code, particularly in hard-to-understand
areas.
- [ ] I have updated the documentation where relevant (API docs, the
reference, and the Sway book).
- [ ] If my change requires substantial documentation changes, I have
[requested support from the DevRel
team](https://github.com/FuelLabs/devrel-requests/issues/new/choose)
- [x] I have added tests that prove my fix is effective or that my
feature works.
- [x] I have added (or requested a maintainer to add) the necessary
`Breaking*` or `New Feature` labels where relevant.
- [x] I have done my best to ensure that my PR adheres to [the Fuel Labs
Code Review
Standards](https://github.com/FuelLabs/rfcs/blob/master/text/code-standards/external-contributors.md).
- [x] I have requested a review from the relevant team or maintainers.

---------

Co-authored-by: kaya <[email protected]>

Author: JoshuaBatty

sway-lsp/benches/lsp_benchmarks/compile.rs

@@ -1,6 +1,7 @@
 use codspeed_criterion_compat::{black_box, criterion_group, Criterion};
 use forc_pkg::manifest::{GenericManifestFile, ManifestFile};
 use lsp_types::Url;
+use parking_lot::RwLock;
 use std::sync::Arc;
 use sway_core::Engines;
 use sway_lsp::core::session;
@@ -9,7 +10,7 @@ use tokio::runtime::Runtime;
 const NUM_DID_CHANGE_ITERATIONS: usize = 10;
 
 fn benchmarks(c: &mut Criterion) {
-    let (uri, session, state) = Runtime::new()
+    let (uri, session, state, _) = Runtime::new()
         .unwrap()
         .block_on(async { black_box(super::compile_test_project().await) });
 
@@ -36,9 +37,11 @@ fn benchmarks(c: &mut Criterion) {
     });
 
     c.bench_function("traverse", |b| {
-        let engines = Engines::default();
-        let results =
-            black_box(session::compile(&build_plan, &engines, None, lsp_mode.as_ref()).unwrap());
+        let engines_original = Arc::new(RwLock::new(Engines::default()));
+        let engines_clone = engines_original.read().clone();
+        let results = black_box(
+            session::compile(&build_plan, &engines_clone, None, lsp_mode.as_ref()).unwrap(),
+        );
         let session = Arc::new(session::Session::new());
         let member_path = sync.member_path(&uri).unwrap();
 
@@ -47,8 +50,10 @@ fn benchmarks(c: &mut Criterion) {
                 session::traverse(
                     member_path.clone(),
                     results.clone(),
-                    &engines,
+                    engines_original.clone(),
+                    &engines_clone,
                     session.clone(),
+                    &state.token_map,
                     lsp_mode.as_ref(),
                 )
                 .unwrap(),

sway-lsp/benches/lsp_benchmarks/mod.rs

@@ -4,36 +4,38 @@ pub mod token_map;
 
 use lsp_types::Url;
 use std::{path::PathBuf, sync::Arc};
+use sway_core::Engines;
 use sway_lsp::{
+    config::GarbageCollectionConfig,
     core::session::{self, Session},
-    server_state::ServerState,
+    server_state::{CompilationContext, ServerState},
 };
 
-pub async fn compile_test_project() -> (Url, Arc<Session>, ServerState) {
+pub async fn compile_test_project() -> (Url, Arc<Session>, ServerState, Engines) {
     // Load the test project
     let uri = Url::from_file_path(benchmark_dir().join("src/main.sw")).unwrap();
     let state = ServerState::default();
+    let engines_clone = state.engines.read().clone();
     let session = Arc::new(Session::new());
     let sync = state.get_or_init_global_sync_workspace(&uri).await.unwrap();
     let temp_uri = sync.workspace_to_temp_url(&uri).unwrap();
 
-    let lsp_mode = Some(sway_core::LspConfig {
+    state.documents.handle_open_file(&temp_uri).await;
+    let ctx = CompilationContext {
+        session: Some(session.clone()),
+        sync: Some(sync.clone()),
+        token_map: state.token_map.clone(),
+        engines: state.engines.clone(),
         optimized_build: false,
         file_versions: Default::default(),
-    });
+        uri: Some(uri.clone()),
+        version: None,
+        gc_options: GarbageCollectionConfig::default(),
+    };
 
-    state.documents.handle_open_file(&temp_uri).await;
     // Compile the project
-    session::parse_project(
-        &temp_uri,
-        &session.engines.read(),
-        None,
-        lsp_mode,
-        session.clone(),
-        &sync,
-    )
-    .unwrap();
-    (temp_uri, session, state)
+    session::parse_project(&temp_uri, &engines_clone, None, &ctx).unwrap();
+    (temp_uri, session, state, engines_clone)
 }
 
 pub fn sway_workspace_dir() -> PathBuf {

sway-lsp/benches/lsp_benchmarks/requests.rs

@@ -3,29 +3,26 @@ use lsp_types::{
     CompletionResponse, DocumentSymbolResponse, Position, Range, TextDocumentContentChangeEvent,
     TextDocumentIdentifier,
 };
-use sway_lsp::{
-    capabilities, config::LspClient, lsp_ext::OnEnterParams, utils::keyword_docs::KeywordDocs,
-};
+use sway_lsp::{capabilities, lsp_ext::OnEnterParams};
 use tokio::runtime::Runtime;
 
 fn benchmarks(c: &mut Criterion) {
-    let (uri, session, state) = Runtime::new()
+    let (uri, session, state, engines) = Runtime::new()
         .unwrap()
         .block_on(async { black_box(super::compile_test_project().await) });
     let sync = state.sync_workspace.get().unwrap();
     let config = sway_lsp::config::Config::default();
-    let keyword_docs = KeywordDocs::new();
     let position = Position::new(1717, 24);
     let range = Range::new(Position::new(1628, 0), Position::new(1728, 0));
 
     c.bench_function("semantic_tokens", |b| {
-        b.iter(|| capabilities::semantic_tokens::semantic_tokens_full(session.clone(), &uri))
+        b.iter(|| capabilities::semantic_tokens::semantic_tokens_full(&state.token_map, &uri))
     });
 
     c.bench_function("document_symbol", |b| {
         b.iter(|| {
             session
-                .document_symbols(&uri)
+                .document_symbols(&uri, &state.token_map, &engines)
                 .map(DocumentSymbolResponse::Nested)
         })
     });
@@ -34,40 +31,44 @@ fn benchmarks(c: &mut Criterion) {
         let position = Position::new(1698, 28);
         b.iter(|| {
             session
-                .completion_items(&uri, position, ".")
+                .completion_items(&uri, position, ".", &state.token_map, &engines)
                 .map(CompletionResponse::Array)
         })
     });
 
     c.bench_function("hover", |b| {
         b.iter(|| {
-            capabilities::hover::hover_data(
+            capabilities::hover::hover_data(&state, &engines, session.clone(), &uri, position)
+        })
+    });
+
+    c.bench_function("highlight", |b| {
+        b.iter(|| {
+            capabilities::highlight::get_highlights(
                 session.clone(),
-                &keyword_docs,
+                &engines,
+                &state.token_map,
                 &uri,
                 position,
-                LspClient::default(),
-                sync,
             )
         })
     });
 
-    c.bench_function("highlight", |b| {
-        b.iter(|| capabilities::highlight::get_highlights(session.clone(), &uri, position))
-    });
-
     c.bench_function("find_all_references", |b| {
-        b.iter(|| session.token_references(&uri, position, sync))
+        b.iter(|| session.token_references(&uri, position, &state.token_map, &engines, sync))
     });
 
     c.bench_function("goto_definition", |b| {
-        b.iter(|| session.token_definition_response(&uri, position, sync))
+        b.iter(|| {
+            session.token_definition_response(&uri, position, &engines, &state.token_map, sync)
+        })
     });
 
     c.bench_function("inlay_hints", |b| {
         b.iter(|| {
             capabilities::inlay_hints::inlay_hints(
-                session.clone(),
+                &engines,
+                &state.token_map,
                 &uri,
                 &range,
                 &config.inlay_hints,
@@ -76,13 +77,16 @@ fn benchmarks(c: &mut Criterion) {
     });
 
     c.bench_function("prepare_rename", |b| {
-        b.iter(|| capabilities::rename::prepare_rename(session.clone(), &uri, position, sync))
+        b.iter(|| {
+            capabilities::rename::prepare_rename(&engines, &state.token_map, &uri, position, sync)
+        })
     });
 
     c.bench_function("rename", |b| {
         b.iter(|| {
             capabilities::rename::rename(
-                session.clone(),
+                &engines,
+                &state.token_map,
                 "new_token_name".to_string(),
                 &uri,
                 position,
@@ -94,7 +98,15 @@ fn benchmarks(c: &mut Criterion) {
     c.bench_function("code_action", |b| {
         let range = Range::new(Position::new(4, 10), Position::new(4, 10));
         b.iter(|| {
-            capabilities::code_actions::code_actions(session.clone(), &range, &uri, &uri, &vec![])
+            capabilities::code_actions::code_actions(
+                session.clone(),
+                &engines,
+                &state.token_map,
+                &range,
+                &uri,
+                &uri,
+                &vec![],
+            )
         })
     });
 

sway-lsp/benches/lsp_benchmarks/token_map.rs

@@ -3,43 +3,49 @@ use lsp_types::Position;
 use tokio::runtime::Runtime;
 
 fn benchmarks(c: &mut Criterion) {
-    let (uri, session, state) = Runtime::new()
+    let (uri, _, state, engines) = Runtime::new()
         .unwrap()
         .block_on(async { black_box(super::compile_test_project().await) });
     let sync = state.sync_workspace.get().unwrap();
-    let engines = session.engines.read();
     let position = Position::new(1716, 24);
+    let path = uri.to_file_path().unwrap();
+    let program_id = sway_lsp::core::session::program_id_from_path(&path, &engines).unwrap();
+    c.bench_function("tokens_for_program", |b| {
+        b.iter(|| {
+            let _: Vec<_> = state.token_map.tokens_for_program(program_id).collect();
+        })
+    });
 
     c.bench_function("tokens_for_file", |b| {
         b.iter(|| {
-            let _: Vec<_> = session.token_map().tokens_for_file(&uri).collect();
+            let _: Vec<_> = state.token_map.tokens_for_file(&uri).collect();
         })
     });
 
     c.bench_function("idents_at_position", |b| {
         b.iter(|| {
-            session
-                .token_map()
-                .idents_at_position(position, session.token_map().iter())
+            state
+                .token_map
+                .idents_at_position(position, state.token_map.iter())
         })
     });
 
     c.bench_function("tokens_at_position", |b| {
         b.iter(|| {
-            session
-                .token_map()
+            state
+                .token_map
                 .tokens_at_position(&engines, &uri, position, None)
         })
     });
 
     c.bench_function("token_at_position", |b| {
-        b.iter(|| session.token_map().token_at_position(&uri, position))
+        b.iter(|| state.token_map.token_at_position(&uri, position))
     });
 
     c.bench_function("parent_decl_at_position", |b| {
         b.iter(|| {
-            session
-                .token_map()
+            state
+                .token_map
                 .parent_decl_at_position(&engines, &uri, position)
         })
     });

sway-lsp/src/capabilities/code_actions/mod.rs

@@ -44,19 +44,19 @@ pub(crate) struct CodeActionContext<'a> {
 
 pub fn code_actions(
     session: Arc<Session>,
+    engines: &Engines,
+    token_map: &TokenMap,
     range: &Range,
     uri: &Url,
     temp_uri: &Url,
     diagnostics: &Vec<Diagnostic>,
 ) -> Option<CodeActionResponse> {
-    let t = session
-        .token_map()
-        .token_at_position(temp_uri, range.start)?;
+    let t = token_map.token_at_position(temp_uri, range.start)?;
     let token = t.value();
 
     let ctx = CodeActionContext {
-        engines: &session.engines.read(),
-        tokens: session.token_map(),
+        engines,
+        tokens: token_map,
         token,
         uri,
         temp_uri,

sway-lsp/src/capabilities/highlight.rs

@@ -1,17 +1,22 @@
-use crate::core::session::Session;
+use crate::core::{session::Session, token_map::TokenMap};
 use lsp_types::{DocumentHighlight, Position, Url};
 use std::sync::Arc;
+use sway_core::Engines;
 
 pub fn get_highlights(
     session: Arc<Session>,
+    engines: &Engines,
+    token_map: &TokenMap,
     url: &Url,
     position: Position,
 ) -> Option<Vec<DocumentHighlight>> {
     let _p = tracing::trace_span!("get_highlights").entered();
-    session.token_ranges(url, position).map(|ranges| {
-        ranges
-            .into_iter()
-            .map(|range| DocumentHighlight { range, kind: None })
-            .collect()
-    })
+    session
+        .token_ranges(engines, token_map, url, position)
+        .map(|ranges| {
+            ranges
+                .into_iter()
+                .map(|range| DocumentHighlight { range, kind: None })
+                .collect()
+        })
 }