mf_search/
step_registry.rs

1use std::any::{Any, TypeId};
2use std::collections::HashMap;
3use std::sync::{Arc, RwLock, OnceLock};
4
5use mf_model::node_pool::NodePool;
6use mf_model::NodeId;
7use mf_transform::step::Step;
8
9use crate::backend::IndexMutation;
10use crate::model::IndexDoc;
11
12/// 步骤转换上下文（提供前后池以便获取节点/父链等）
13pub struct StepIndexContext<'a> {
14    pub pool_before: &'a NodePool,
15    pub pool_after: &'a NodePool,
16}
17
18/// 类型安全的 Step→索引增量 转换器
19pub trait TypedStepIndexer<T>: Send + Sync + 'static
20where
21    T: Step + 'static,
22{
23    fn index_step(
24        &self,
25        step: &T,
26        ctx: &StepIndexContext,
27    ) -> Vec<IndexMutation>;
28    fn name() -> &'static str
29    where
30        Self: Sized,
31    {
32        std::any::type_name::<Self>()
33    }
34}
35
36/// 类型擦除的转换器
37struct ErasedIndexer {
38    type_id: TypeId,
39    index_fn: fn(&dyn Any, &StepIndexContext) -> Vec<IndexMutation>,
40}
41
42impl ErasedIndexer {
43    fn new<T, C>() -> Self
44    where
45        T: Step + 'static,
46        C: TypedStepIndexer<T> + Default + 'static,
47    {
48        Self {
49            type_id: TypeId::of::<T>(),
50            index_fn: |step_any, ctx| {
51                let converter = C::default();
52                let step = step_any
53                    .downcast_ref::<T>()
54                    .expect("Type mismatch in step indexer");
55                converter.index_step(step, ctx)
56            },
57        }
58    }
59
60    fn try_index(
61        &self,
62        step: &dyn Step,
63        ctx: &StepIndexContext,
64    ) -> Option<Vec<IndexMutation>> {
65        if step.type_id() != self.type_id {
66            return None;
67        }
68        Some((self.index_fn)(step as &dyn Any, ctx))
69    }
70}
71
72/// 转换器注册表
73#[derive(Default)]
74pub struct StepIndexerRegistry {
75    by_type: HashMap<TypeId, Arc<ErasedIndexer>>,
76}
77
78impl StepIndexerRegistry {
79    pub fn new() -> Self {
80        Self { by_type: HashMap::new() }
81    }
82
83    pub fn register<T, C>(&mut self) -> &mut Self
84    where
85        T: Step + 'static,
86        C: TypedStepIndexer<T> + Default + 'static,
87    {
88        let type_id = TypeId::of::<T>();
89        self.by_type
90            .entry(type_id)
91            .or_insert_with(|| Arc::new(ErasedIndexer::new::<T, C>()));
92        self
93    }
94
95    pub fn index_step(
96        &self,
97        step: &dyn Step,
98        ctx: &StepIndexContext,
99    ) -> Option<Vec<IndexMutation>> {
100        let type_id = step.type_id();
101        self.by_type.get(&type_id).and_then(|e| e.try_index(step, ctx))
102    }
103}
104
105static GLOBAL: OnceLock<RwLock<StepIndexerRegistry>> = OnceLock::new();
106
107pub fn global_registry() -> &'static RwLock<StepIndexerRegistry> {
108    GLOBAL.get_or_init(|| RwLock::new(StepIndexerRegistry::new()))
109}
110
111pub fn register_step_indexer<T, C>()
112where
113    T: Step + 'static,
114    C: TypedStepIndexer<T> + Default + 'static,
115{
116    let mut reg = global_registry().write().unwrap();
117    reg.register::<T, C>();
118}
119
120// ---------------- 内置步骤的默认转换器 ----------------
121use mf_transform::attr_step::AttrStep;
122use mf_transform::mark_step::{AddMarkStep, RemoveMarkStep};
123use mf_transform::node_step::{AddNodeStep, RemoveNodeStep, MoveNodeStep};
124use mf_model::node_definition::NodeTree;
125use serde::Deserialize;
126
127#[derive(Default)]
128struct AttrIndexer;
129impl TypedStepIndexer<AttrStep> for AttrIndexer {
130    fn index_step(
131        &self,
132        step: &AttrStep,
133        ctx: &StepIndexContext,
134    ) -> Vec<IndexMutation> {
135        if let Some(node) = ctx.pool_after.get_node(&step.id) {
136            vec![IndexMutation::Upsert(IndexDoc::from_node(
137                ctx.pool_after,
138                &node,
139            ))]
140        } else {
141            Vec::new()
142        }
143    }
144}
145
146#[derive(Default)]
147struct AddMarkIndexer;
148impl TypedStepIndexer<AddMarkStep> for AddMarkIndexer {
149    fn index_step(
150        &self,
151        step: &AddMarkStep,
152        ctx: &StepIndexContext,
153    ) -> Vec<IndexMutation> {
154        if let Some(node) = ctx.pool_after.get_node(&step.id) {
155            vec![IndexMutation::Upsert(IndexDoc::from_node(
156                ctx.pool_after,
157                &node,
158            ))]
159        } else {
160            Vec::new()
161        }
162    }
163}
164
165#[derive(Default)]
166struct RemoveMarkIndexer;
167impl TypedStepIndexer<RemoveMarkStep> for RemoveMarkIndexer {
168    fn index_step(
169        &self,
170        step: &RemoveMarkStep,
171        ctx: &StepIndexContext,
172    ) -> Vec<IndexMutation> {
173        if let Some(node) = ctx.pool_after.get_node(&step.id) {
174            vec![IndexMutation::Upsert(IndexDoc::from_node(
175                ctx.pool_after,
176                &node,
177            ))]
178        } else {
179            Vec::new()
180        }
181    }
182}
183
184#[derive(Default)]
185struct AddNodeIndexer;
186impl TypedStepIndexer<AddNodeStep> for AddNodeIndexer {
187    fn index_step(
188        &self,
189        step: &AddNodeStep,
190        ctx: &StepIndexContext,
191    ) -> Vec<IndexMutation> {
192        let mut muts = Vec::new();
193        for ne in &step.nodes {
194            collect_adds_for_node_enum(ctx.pool_after, ne, &mut muts);
195        }
196        muts
197    }
198}
199
200#[derive(Default)]
201struct RemoveNodeIndexer;
202impl TypedStepIndexer<RemoveNodeStep> for RemoveNodeIndexer {
203    fn index_step(
204        &self,
205        step: &RemoveNodeStep,
206        ctx: &StepIndexContext,
207    ) -> Vec<IndexMutation> {
208        let mut all_ids: Vec<String> = Vec::new();
209        for id in &step.node_ids {
210            if let Some(enum_subtree) =
211                ctx.pool_before.get_inner().all_children(id, None)
212            {
213                all_ids.extend(
214                    collect_ids_from_enum(&enum_subtree)
215                        .into_iter()
216                        .map(|id| id.to_string()),
217                );
218            } else {
219                all_ids.push(id.to_string());
220            }
221        }
222        vec![IndexMutation::DeleteManyById(all_ids)]
223    }
224}
225
226#[derive(Default)]
227struct MoveNodeIndexer;
228impl TypedStepIndexer<MoveNodeStep> for MoveNodeIndexer {
229    fn index_step(
230        &self,
231        step: &MoveNodeStep,
232        ctx: &StepIndexContext,
233    ) -> Vec<IndexMutation> {
234        // MoveNodeStep 的字段为私有，使用序列化提取 node_id
235        #[derive(Deserialize)]
236        struct MoveNodeSerde {
237            node_id: NodeId,
238        }
239        if let Some(bytes) = Step::serialize(step) {
240            if let Ok(ms) = serde_json::from_slice::<MoveNodeSerde>(&bytes) {
241                let mut muts = Vec::new();
242                if let Some(enum_subtree) =
243                    ctx.pool_after.get_inner().all_children(&ms.node_id, None)
244                {
245                    collect_upserts_for_enum(
246                        ctx.pool_after,
247                        &enum_subtree,
248                        &mut muts,
249                    );
250                } else if let Some(node) = ctx.pool_after.get_node(&ms.node_id)
251                {
252                    muts.push(IndexMutation::Upsert(IndexDoc::from_node(
253                        ctx.pool_after,
254                        &node,
255                    )));
256                }
257                return muts;
258            }
259        }
260        Vec::new()
261    }
262}
263
264fn collect_ids_from_enum(ne: &NodeTree) -> Vec<NodeId> {
265    let mut ids = vec![ne.0.id.clone()];
266    for c in &ne.1 {
267        ids.extend(collect_ids_from_enum(c));
268    }
269    ids
270}
271
272fn collect_adds_for_node_enum(
273    pool: &NodePool,
274    ne: &NodeTree,
275    out: &mut Vec<IndexMutation>,
276) {
277    let node = std::sync::Arc::new(ne.0.clone());
278    out.push(IndexMutation::Add(IndexDoc::from_node(pool, &node)));
279    for c in &ne.1 {
280        collect_adds_for_node_enum(pool, c, out);
281    }
282}
283
284fn collect_upserts_for_enum(
285    pool: &NodePool,
286    ne: &NodeTree,
287    out: &mut Vec<IndexMutation>,
288) {
289    let node = std::sync::Arc::new(ne.0.clone());
290    out.push(IndexMutation::Upsert(IndexDoc::from_node(pool, &node)));
291    for c in &ne.1 {
292        collect_upserts_for_enum(pool, c, out);
293    }
294}
295
296use std::sync::OnceLock as _OnceLock;
297static DEFAULTS: _OnceLock<()> = _OnceLock::new();
298
299pub fn ensure_default_step_indexers() {
300    DEFAULTS.get_or_init(|| {
301        register_step_indexer::<AttrStep, AttrIndexer>();
302        register_step_indexer::<AddMarkStep, AddMarkIndexer>();
303        register_step_indexer::<RemoveMarkStep, RemoveMarkIndexer>();
304        register_step_indexer::<AddNodeStep, AddNodeIndexer>();
305        register_step_indexer::<RemoveNodeStep, RemoveNodeIndexer>();
306        register_step_indexer::<MoveNodeStep, MoveNodeIndexer>();
307    });
308}
mf_search/step_registry.rs

mf_search/
step_registry.rs
