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}