1use std::{borrow::Borrow, marker::PhantomData, num::NonZeroU16, path::Path};
2
3use rpds::RedBlackTreeMapSync;
4use tinymist_std::ImmutPath;
5use typst::diag::FileResult;
6
7use crate::{AccessModel, Bytes, FileId, FileSnapshot, PathAccessModel};
8
9pub(crate) type RawFileId = NonZeroU16;
10
11#[derive(Default, Debug, Clone)]
14pub struct OverlayAccessModel<K, M, S = K>
15where
16 S: Ord,
17{
18 files: RedBlackTreeMapSync<S, FileSnapshot>,
19 pub inner: M,
21 _key: PhantomData<fn() -> K>,
22}
23
24impl<K, M, S> OverlayAccessModel<K, M, S>
25where
26 S: Ord,
27{
28 pub fn new(inner: M) -> Self {
30 Self {
31 files: RedBlackTreeMapSync::default(),
32 inner,
33 _key: PhantomData,
34 }
35 }
36
37 pub fn inner(&self) -> &M {
39 &self.inner
40 }
41
42 pub fn inner_mut(&mut self) -> &mut M {
44 &mut self.inner
45 }
46
47 pub fn clear_shadow(&mut self) {
49 self.files = RedBlackTreeMapSync::default();
50 }
51}
52
53impl<M> OverlayAccessModel<ImmutPath, M> {
54 pub fn file_paths(&self) -> Vec<ImmutPath> {
56 self.files.keys().cloned().collect()
57 }
58
59 pub fn add_file<Q: Ord + ?Sized>(
61 &mut self,
62 path: &Q,
63 snap: FileSnapshot,
64 cast: impl Fn(&Q) -> ImmutPath,
65 ) where
66 ImmutPath: Borrow<Q>,
67 {
68 match self.files.get_mut(path) {
69 Some(e) => {
70 *e = snap;
71 }
72 None => {
73 self.files.insert_mut(cast(path), snap);
74 }
75 }
76 }
77
78 pub fn remove_file<Q: Ord + ?Sized>(&mut self, path: &Q)
80 where
81 ImmutPath: Borrow<Q>,
82 {
83 self.files.remove_mut(path);
84 }
85}
86
87impl<M> OverlayAccessModel<FileId, M, RawFileId> {
88 pub fn file_paths(&self) -> Vec<FileId> {
90 self.files.keys().copied().map(FileId::from_raw).collect()
91 }
92
93 pub fn add_file(&mut self, id: &FileId, snap: FileSnapshot, cast: impl Fn(&FileId) -> FileId) {
95 match self.files.get_mut(&id.into_raw()) {
96 Some(e) => {
97 *e = snap;
98 }
99 None => {
100 self.files.insert_mut(cast(id).into_raw(), snap);
101 }
102 }
103 }
104
105 pub fn remove_file(&mut self, id: &FileId) {
107 self.files.remove_mut(&id.into_raw());
108 }
109}
110
111impl<M: PathAccessModel> PathAccessModel for OverlayAccessModel<ImmutPath, M> {
112 fn content(&self, src: &Path) -> FileResult<Bytes> {
113 if let Some(content) = self.files.get(src) {
114 return content.content().cloned();
115 }
116
117 self.inner.content(src)
118 }
119}
120
121impl<M: AccessModel> AccessModel for OverlayAccessModel<FileId, M, RawFileId> {
122 fn reset(&mut self) {
123 self.inner.reset();
124 }
125
126 fn content(&self, src: FileId) -> (Option<ImmutPath>, FileResult<Bytes>) {
127 if let Some(content) = self.files.get(&src.into_raw()) {
128 return (None, content.content().cloned());
129 }
130
131 self.inner.content(src)
132 }
133}