markdown_ppp/ast_specialized/
utilities.rs

1//! Utility functions for working with specialized AST types
2//!
3//! This module provides helper functions and utilities for common operations
4//! with specialized AST types.
5
6use super::element_id::IdGenerator;
7use super::type_aliases::with_ids;
8use super::ElementId;
9
10/// Utility functions for adding IDs to AST nodes
11pub mod id_utils {
12    use super::*;
13    use crate::ast::convert::WithData;
14    use crate::ast::generic;
15    use crate::ast::map_data_visitor::{map_user_data, MapDataVisitor};
16
17    /// Add sequential IDs to all nodes in a document
18    pub fn add_ids_to_document(doc: crate::ast::Document) -> with_ids::Document {
19        let doc_with_unit: generic::Document<()> = doc.with_default_data();
20        let mut id_gen = IdGenerator::new();
21        map_user_data(doc_with_unit, |_| id_gen.generate())
22    }
23
24    /// Add sequential IDs starting from a specific value
25    pub fn add_ids_from(doc: crate::ast::Document, start_id: u64) -> with_ids::Document {
26        let doc_with_unit: generic::Document<()> = doc.with_default_data();
27        let mut id_gen = IdGenerator::starting_from(start_id);
28        map_user_data(doc_with_unit, |_| id_gen.generate())
29    }
30
31    /// Custom visitor for adding IDs with more control
32    pub struct IdAssignmentVisitor {
33        id_gen: IdGenerator,
34    }
35
36    impl IdAssignmentVisitor {
37        pub fn new() -> Self {
38            Self {
39                id_gen: IdGenerator::new(),
40            }
41        }
42
43        pub fn starting_from(start_id: u64) -> Self {
44            Self {
45                id_gen: IdGenerator::starting_from(start_id),
46            }
47        }
48    }
49
50    impl Default for IdAssignmentVisitor {
51        fn default() -> Self {
52            Self::new()
53        }
54    }
55
56    impl<T> MapDataVisitor<T, ElementId> for IdAssignmentVisitor {
57        fn map_data(&mut self, _data: T) -> ElementId {
58            self.id_gen.generate()
59        }
60    }
61
62    /// Add IDs to any generic document
63    pub fn add_ids_to_generic_document<T>(doc: generic::Document<T>) -> with_ids::Document {
64        let mut visitor = IdAssignmentVisitor::new();
65        visitor.visit_document(doc)
66    }
67
68    #[cfg(test)]
69    mod tests {
70        use super::*;
71        use crate::ast::{Block, Heading, HeadingKind, Inline};
72
73        #[test]
74        fn test_add_ids_to_document() {
75            let doc = crate::ast::Document {
76                blocks: vec![Block::Heading(Heading {
77                    kind: HeadingKind::Atx(1),
78                    content: vec![Inline::Text("Test".to_string())],
79                })],
80            };
81
82            let doc_with_ids = add_ids_to_document(doc);
83
84            // Should have ID assigned to document
85            assert!(doc_with_ids.user_data.id() > 0);
86
87            // Should have IDs assigned to all blocks
88            assert_eq!(doc_with_ids.blocks.len(), 1);
89            match &doc_with_ids.blocks[0] {
90                generic::Block::Heading(h) => {
91                    assert!(h.user_data.id() > 0);
92                    assert_eq!(h.content.len(), 1);
93                    match &h.content[0] {
94                        generic::Inline::Text { user_data, .. } => {
95                            assert!(user_data.id() > 0);
96                        }
97                        _ => panic!("Expected text inline"),
98                    }
99                }
100                _ => panic!("Expected heading"),
101            }
102        }
103
104        #[test]
105        fn test_add_ids_from() {
106            let doc = crate::ast::Document {
107                blocks: vec![Block::Heading(Heading {
108                    kind: HeadingKind::Atx(1),
109                    content: vec![Inline::Text("Test".to_string())],
110                })],
111            };
112
113            let doc_with_ids = add_ids_from(doc, 100);
114
115            // Should start from the specified ID
116            assert!(doc_with_ids.user_data.id() >= 100);
117        }
118
119        #[test]
120        fn test_id_assignment_visitor() {
121            let mut visitor = IdAssignmentVisitor::starting_from(50);
122            let id1 = visitor.map_data(());
123            let id2 = visitor.map_data(());
124
125            assert_eq!(id1.id(), 50);
126            assert_eq!(id2.id(), 51);
127        }
128    }
129}