1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

use pax_manifest::{MappedString, Token};
use pax_runtime::api::serde::{Deserialize, Serialize};
use regex::Regex;
use std::io::Write;
use std::{
    collections::{BTreeMap, HashMap},
    fs::{File, OpenOptions},
    io::{BufRead, BufReader},
};

#[derive(Debug, Serialize, Deserialize)]
#[serde(crate = "pax_runtime::api::serde")]
pub struct RangeData {
    pub end: usize,
    pub id: usize,
    pub token: Token,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(crate = "pax_runtime::api::serde")]
pub struct SourceMap {
    pub sources: HashMap<usize, Token>,
    pub ranges: BTreeMap<usize, RangeData>,
    pub next_id: usize,
}

impl SourceMap {
    pub fn new() -> Self {
        SourceMap {
            sources: HashMap::new(),
            ranges: BTreeMap::new(),
            next_id: 0,
        }
    }

    pub fn generate_id(&mut self) -> usize {
        let id = self.next_id;
        self.next_id += 1;
        id
    }

    pub fn insert(&mut self, token: Token) -> usize {
        let id = self.generate_id();
        self.sources.insert(id, token);
        id
    }

    pub fn generate_start_end_marker(&mut self, id: usize) -> (String, String) {
        let start_marker = format!("/* source_map_start_{} */", id);
        let end_marker = format!("/* source_map_end_{} */", id);
        (start_marker, end_marker)
    }

    pub fn generate_mapped_string(&mut self, content: String, id: usize) -> MappedString {
        let (start_marker, end_marker) = self.generate_start_end_marker(id);
        MappedString {
            content,
            source_map_start_marker: Some(start_marker),
            source_map_end_marker: Some(end_marker),
        }
    }

    pub fn extract_ranges_from_generated_code(&mut self, file_path: &str) {
        let file = File::open(file_path).expect("Failed to open file");
        let reader = BufReader::new(file);

        let start_regex = Regex::new(r"/\* source_map_start_(\d+) \*/").unwrap();
        let end_regex = Regex::new(r"/\* source_map_end_(\d+) \*/").unwrap();

        let mut start_positions: HashMap<usize, usize> = HashMap::new();
        let mut processed_content = String::new();

        for (line_num, res) in reader.lines().enumerate() {
            let line = res.expect("Invalid line");
            let mut line_processed = line.clone();

            if let Some(captures) = start_regex.captures(&line) {
                let id: usize = captures[1].parse().unwrap();
                let start_pos = line_num + 1;
                start_positions.insert(id, start_pos);
                line_processed = line_processed.replace(&captures[0], "");
            }

            if let Some(captures) = end_regex.captures(&line) {
                let id: usize = captures[1].parse().unwrap();
                if let Some(start_pos) = start_positions.remove(&id) {
                    let end_pos = line_num + 1;
                    let token = self
                        .sources
                        .get(&id)
                        .cloned()
                        .unwrap_or_else(|| Token::default());
                    let range_data = RangeData {
                        end: end_pos,
                        id,
                        token,
                    };
                    self.ranges.insert(start_pos, range_data);
                }
                line_processed = line_processed.replace(&captures[0], "");
            }

            processed_content.push_str(&line_processed);
            processed_content.push('\n');
        }

        let mut file = OpenOptions::new()
            .write(true)
            .truncate(true)
            .open(file_path)
            .expect("Failed to open file in write mode");
        write!(file, "{}", processed_content).expect("Failed to write to file");
    }

    pub fn get_range_for_line(&self, line: usize) -> Option<&RangeData> {
        let (_, range) = self.ranges.range(..=line).next_back()?;
        if line <= range.end {
            Some(range)
        } else {
            None
        }
    }
}