Skip to main content

docparser_core/
registry.rs

1//! 注册表:能力 / 后端注册、按策略解析后端链、复合后端委托环检测。
2
3use std::collections::HashMap;
4use std::fmt;
5use std::sync::Arc;
6
7use crate::plugin::{FormatPlugin, ParseCapability, ProcessCapability};
8use crate::policy::BackendPolicy;
9
10/// 复合后端不可委托成环。
11#[derive(Debug)]
12pub struct RegistryError(pub String);
13
14impl fmt::Display for RegistryError {
15    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
16        write!(f, "registry: {}", self.0)
17    }
18}
19
20impl std::error::Error for RegistryError {}
21
22#[derive(Default)]
23pub struct Registry {
24    parse_caps: HashMap<String, Vec<Arc<dyn ParseCapability>>>,
25    process_caps: HashMap<String, Arc<dyn ProcessCapability>>,
26    plugins: Vec<Arc<dyn FormatPlugin>>,
27    delegates: HashMap<String, String>, // backend id -> delegate id
28}
29
30impl Registry {
31    pub fn new() -> Self {
32        Self::default()
33    }
34
35    /// 注册解析后端;同一声明的多个 format 共享同一 `Arc`。
36    pub fn register_parse(
37        &mut self,
38        cap: Arc<dyn ParseCapability>,
39        delegate_to: Option<&str>,
40    ) -> Result<(), RegistryError> {
41        if let Some(dst) = delegate_to {
42            self.check_delegate_cycle(cap.id(), dst)?;
43            self.delegates.insert(cap.id().to_string(), dst.to_string());
44        }
45        for f in cap.formats() {
46            self.parse_caps
47                .entry((*f).to_string())
48                .or_default()
49                .push(cap.clone());
50        }
51        Ok(())
52    }
53
54    pub fn register_process(&mut self, cap: Arc<dyn ProcessCapability>) {
55        self.process_caps.insert(cap.id().to_string(), cap);
56    }
57
58    pub fn register_plugin(&mut self, plugin: Arc<dyn FormatPlugin>) {
59        self.plugins.push(plugin);
60    }
61
62    /// 按格式找插件。
63    pub fn plugin_for(&self, fmt: &str) -> Option<Arc<dyn FormatPlugin>> {
64        self.plugins
65            .iter()
66            .find(|p| p.formats().contains(&fmt))
67            .cloned()
68    }
69
70    fn check_delegate_cycle(&self, src: &str, dst: &str) -> Result<(), RegistryError> {
71        if src == dst {
72            return Err(RegistryError(format!(
73                "backend cannot delegate to itself: {src}"
74            )));
75        }
76        let mut seen = std::collections::HashSet::new();
77        let mut cur = Some(dst.to_string());
78        while let Some(c) = cur {
79            if c == src {
80                return Err(RegistryError(format!(
81                    "delegate cycle: {src} -> ... -> {dst}"
82                )));
83            }
84            if !seen.insert(c.clone()) {
85                break;
86            }
87            cur = self.delegates.get(&c).cloned();
88        }
89        Ok(())
90    }
91
92    /// 解析该格式可尝试的解析后端链(过滤不可用项)。
93    pub fn resolve_parse_chain(
94        &self,
95        fmt: &str,
96        policy: &BackendPolicy,
97        default: &[&str],
98    ) -> Vec<Arc<dyn ParseCapability>> {
99        let chain =
100            policy.ordered_chain(&default.iter().map(|s| s.to_string()).collect::<Vec<_>>());
101        let caps = self.parse_caps.get(fmt).cloned().unwrap_or_default();
102        chain
103            .iter()
104            .filter_map(|bid| {
105                caps.iter()
106                    .find(|c| c.id() == bid)
107                    .filter(|c| c.is_available())
108                    .cloned()
109            })
110            .collect()
111    }
112
113    pub fn resolve_process_chain(
114        &self,
115        policy: &BackendPolicy,
116        default: &[&str],
117    ) -> Vec<Arc<dyn ProcessCapability>> {
118        let chain =
119            policy.ordered_chain(&default.iter().map(|s| s.to_string()).collect::<Vec<_>>());
120        chain
121            .iter()
122            .filter_map(|bid| self.process_caps.get(bid.as_str()).cloned())
123            .filter(|c| c.is_available())
124            .filter(|c| c.id() != "ocr" || policy.enable_ocr)
125            .collect()
126    }
127
128    pub fn list_backends(&self) -> Vec<String> {
129        let mut ids: Vec<String> = self
130            .parse_caps
131            .values()
132            .flatten()
133            .map(|c| format!("{} ({})", c.id(), c.kind()))
134            .chain(self.process_caps.keys().cloned())
135            .collect();
136        ids.sort();
137        ids.dedup();
138        ids
139    }
140}
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145    use crate::error::Result;
146    use crate::plugin::RawDoc;
147
148    struct MockParse {
149        id: &'static str,
150        formats: Vec<&'static str>,
151        available: bool,
152    }
153    impl ParseCapability for MockParse {
154        fn id(&self) -> &str {
155            self.id
156        }
157        fn formats(&self) -> &[&str] {
158            &self.formats
159        }
160        fn is_available(&self) -> bool {
161            self.available
162        }
163        fn parse(&self, _path: &str) -> Result<RawDoc> {
164            Ok(RawDoc::default())
165        }
166    }
167
168    struct MockProc {
169        id: &'static str,
170    }
171    impl ProcessCapability for MockProc {
172        fn id(&self) -> &str {
173            self.id
174        }
175        fn process(&self, _raw: &mut RawDoc) -> Result<()> {
176            Ok(())
177        }
178    }
179
180    #[test]
181    fn register_and_resolve_multi_format() {
182        let mut reg = Registry::new();
183        let cap: Arc<dyn ParseCapability> = Arc::new(MockParse {
184            id: "fmt",
185            formats: vec!["a", "b"],
186            available: true,
187        });
188        reg.register_parse(cap, None).unwrap();
189        let pol = BackendPolicy::default();
190        assert_eq!(reg.resolve_parse_chain("a", &pol, &["fmt"]).len(), 1);
191        assert_eq!(reg.resolve_parse_chain("b", &pol, &["fmt"]).len(), 1);
192        assert_eq!(reg.resolve_parse_chain("c", &pol, &["fmt"]).len(), 0);
193    }
194
195    #[test]
196    fn unavailable_is_skipped() {
197        let mut reg = Registry::new();
198        reg.register_parse(
199            Arc::new(MockParse {
200                id: "x",
201                formats: vec!["a"],
202                available: false,
203            }),
204            None,
205        )
206        .unwrap();
207        let pol = BackendPolicy::default();
208        assert_eq!(reg.resolve_parse_chain("a", &pol, &["x"]).len(), 0);
209    }
210
211    #[test]
212    fn process_chain_respects_enable_ocr() {
213        let mut reg = Registry::new();
214        reg.register_process(Arc::new(MockProc { id: "ocr" }));
215        let on = BackendPolicy {
216            enable_ocr: true,
217            ..Default::default()
218        };
219        let off = BackendPolicy {
220            enable_ocr: false,
221            ..Default::default()
222        };
223        assert_eq!(reg.resolve_process_chain(&on, &["ocr"]).len(), 1);
224        assert_eq!(reg.resolve_process_chain(&off, &["ocr"]).len(), 0);
225    }
226
227    #[test]
228    fn delegate_cycle_detected() {
229        let mut reg = Registry::new();
230        reg.register_parse(
231            Arc::new(MockParse {
232                id: "a",
233                formats: vec!["x"],
234                available: true,
235            }),
236            Some("b"),
237        )
238        .unwrap();
239        let err = reg.register_parse(
240            Arc::new(MockParse {
241                id: "b",
242                formats: vec!["x"],
243                available: true,
244            }),
245            Some("a"),
246        );
247        assert!(err.is_err());
248        // 自环同样拒绝
249        let err2 = reg.register_parse(
250            Arc::new(MockParse {
251                id: "c",
252                formats: vec!["x"],
253                available: true,
254            }),
255            Some("c"),
256        );
257        assert!(err2.is_err());
258    }
259}