jsona_lsp/
world.rs

1use crate::{
2    config::{InitializationOptions, LspConfig},
3    lsp_ext::notification::{InitializeWorkspace, InitializeWorkspaceParams},
4};
5use arc_swap::ArcSwap;
6use jsona::{
7    dom::{Keys, Node},
8    parser::Parse,
9};
10use jsona_schema::Schema;
11use jsona_util::{
12    environment::Environment,
13    schema::{
14        associations::{priority, source, AssociationRule, SchemaAssociation},
15        Schemas,
16    },
17    AsyncRwLock, HashMap, IndexMap,
18};
19use lsp_async_stub::{rpc, util::Mapper, Context, RequestWriter};
20use lsp_types::Url;
21use once_cell::sync::Lazy;
22use serde_json::{json, Value};
23use std::sync::Arc;
24pub type World<E> = Arc<WorldState<E>>;
25
26#[repr(transparent)]
27pub struct Workspaces<E: Environment>(IndexMap<Url, WorkspaceState<E>>);
28
29impl<E: Environment> std::ops::Deref for Workspaces<E> {
30    type Target = IndexMap<Url, WorkspaceState<E>>;
31
32    fn deref(&self) -> &Self::Target {
33        &self.0
34    }
35}
36
37impl<E: Environment> std::ops::DerefMut for Workspaces<E> {
38    fn deref_mut(&mut self) -> &mut Self::Target {
39        &mut self.0
40    }
41}
42
43impl<E: Environment> Workspaces<E> {
44    pub fn by_document(&self, document_uri: &Url) -> &WorkspaceState<E> {
45        self.find_workspace(document_uri).unwrap().1
46    }
47
48    pub fn by_document_mut(&mut self, document_uri: &Url) -> &mut WorkspaceState<E> {
49        let uri = self.find_workspace(document_uri).unwrap().0.clone();
50        self.0.get_mut(&uri).unwrap()
51    }
52
53    pub fn try_get_document(
54        &self,
55        document_uri: &Url,
56    ) -> Result<(&WorkspaceState<E>, &DocumentState), rpc::Error> {
57        let ws = self.by_document(document_uri);
58        let doc = ws.try_get_document(document_uri)?;
59        Ok((ws, doc))
60    }
61
62    fn find_workspace(&self, document_uri: &Url) -> Option<(&Url, &WorkspaceState<E>)> {
63        self.0
64            .iter()
65            .filter(|(key, _)| {
66                document_uri.as_str().starts_with(key.as_str()) || *key == &*DEFAULT_WORKSPACE_URI
67            })
68            .max_by(|(a, _), (b, _)| a.as_str().len().cmp(&b.as_str().len()))
69            .or_else(|| self.0.first())
70    }
71}
72
73pub struct WorldState<E: Environment> {
74    pub(crate) env: E,
75    pub(crate) id: String,
76    pub(crate) workspaces: AsyncRwLock<Workspaces<E>>,
77    pub(crate) initialization_options: ArcSwap<InitializationOptions>,
78}
79
80pub static DEFAULT_WORKSPACE_URI: Lazy<Url> = Lazy::new(|| Url::parse("root:///").unwrap());
81
82impl<E: Environment> WorldState<E> {
83    pub fn new(env: E) -> Self {
84        let id = format!(
85            "{:x}",
86            md5::compute(format!("JSONA-{}", env.now().unix_timestamp_nanos()))
87        );
88        Self {
89            id,
90            workspaces: AsyncRwLock::new(Workspaces(IndexMap::default())),
91            initialization_options: Default::default(),
92            env,
93        }
94    }
95}
96
97pub struct WorkspaceState<E: Environment> {
98    pub(crate) root: Url,
99    pub(crate) documents: HashMap<lsp_types::Url, DocumentState>,
100    pub(crate) schemas: Schemas<E>,
101    pub(crate) lsp_config: LspConfig,
102}
103
104impl<E: Environment> WorkspaceState<E> {
105    pub(crate) fn new(env: E, root: Url) -> Self {
106        Self {
107            root,
108            documents: Default::default(),
109            schemas: Schemas::new(env),
110            lsp_config: LspConfig::default(),
111        }
112    }
113}
114
115impl<E: Environment> WorkspaceState<E> {
116    pub(crate) fn try_get_document(
117        &self,
118        document_uri: &Url,
119    ) -> Result<&DocumentState, rpc::Error> {
120        self.documents.get(document_uri).ok_or_else(|| {
121            tracing::debug!(%document_uri, "not found document in workspace");
122            rpc::Error::invalid_params()
123        })
124    }
125
126    #[tracing::instrument(skip_all, fields(%self.root))]
127    pub(crate) async fn initialize(
128        &mut self,
129        context: Context<World<E>>,
130        lsp_config: &Value,
131    ) -> Result<(), anyhow::Error> {
132        if let Err(error) = self.lsp_config.update_from_json(lsp_config) {
133            tracing::error!(?error, "invalid configuration");
134        }
135
136        if self.lsp_config.schema.cache {
137            let cache_path = context.initialization_options.load().cache_path.clone();
138            self.schemas.set_cache_path(cache_path);
139        } else {
140            self.schemas.set_cache_path(None);
141        }
142
143        self.schemas.associations().clear();
144
145        if !self.lsp_config.schema.enabled {
146            return Ok(());
147        }
148
149        let store_url = self.lsp_config.schema.store_url.clone();
150
151        if let Err(error) = self
152            .schemas
153            .associations()
154            .add_from_schemastore(&store_url, &Some(self.root.clone()))
155            .await
156        {
157            let store_url = store_url.as_ref().map(|v| v.as_str());
158            tracing::error!(%error, ?store_url, "failed to load schemastore");
159        }
160
161        for (name, items) in &self.lsp_config.schema.associations {
162            match self.schemas.associations().get_schema_url(name) {
163                Some(schema_uri) => {
164                    let assoc = SchemaAssociation {
165                        url: schema_uri.clone(),
166                        meta: json!({
167                            "source": source::LSP_CONFIG,
168                        }),
169                        priority: priority::LSP_CONFIG,
170                    };
171                    match AssociationRule::batch(items, &Some(self.root.clone())) {
172                        Ok(rules) => {
173                            for rule in rules {
174                                self.schemas.associations().add(rule, assoc.clone())
175                            }
176                        }
177                        Err(error) => {
178                            tracing::error!(%error, %schema_uri, "failed to add schema associations");
179                        }
180                    }
181                }
182                None => {
183                    tracing::error!(%name, "failed to add schema associations");
184                }
185            }
186        }
187
188        self.refresh_associated_schemas().await;
189
190        self.emit_initialize_workspace(context.clone()).await;
191
192        Ok(())
193    }
194
195    pub(crate) async fn emit_initialize_workspace(&self, mut context: Context<World<E>>) {
196        if let Err(error) = context
197            .write_notification::<InitializeWorkspace, _>(Some(InitializeWorkspaceParams {
198                root_uri: self.root.clone(),
199            }))
200            .await
201        {
202            tracing::error!(%error, "failed to write notification");
203        }
204    }
205
206    pub(crate) async fn query_schemas(&self, file: &Url, path: &Keys) -> Option<Vec<Schema>> {
207        let schema_association = self.schemas.associations().query_for(file)?;
208        match self.schemas.query(&schema_association.url, path).await {
209            Ok(v) => Some(v),
210            Err(error) => {
211                tracing::error!(?error, "failed to query schemas");
212                None
213            }
214        }
215    }
216
217    pub(crate) async fn refresh_associated_schemas(&self) {
218        for (document_uri, doc) in self.documents.iter() {
219            let association = self.schemas.associations().query_for(document_uri);
220            if association.is_none() {
221                self.schemas
222                    .associations()
223                    .add_from_document(document_uri, &doc.dom);
224            }
225        }
226    }
227}
228
229#[derive(Debug, Clone)]
230pub struct DocumentState {
231    pub(crate) parse: Parse,
232    pub(crate) dom: Node,
233    pub(crate) mapper: Mapper,
234}