1use std::{fs::File, path::Path};
2
3use fwob_core::{Key, KeyType, OwnedFrame, Schema};
4
5use crate::{Reader, Result, V1Error, Writer, WriterOptions};
6
7pub struct InMemoryEditor {
14 schema: Schema,
15 title: String,
16 string_table: Vec<String>,
17 frames: Vec<OwnedFrame>,
18 key_type: KeyType,
19}
20
21impl InMemoryEditor {
22 pub fn open(path: impl AsRef<Path>, key_field_index: usize) -> Result<Self> {
23 let mut reader = Reader::open(path, key_field_index)?;
24 reader.verify_key_order()?;
25 let schema = reader.schema().clone();
26 let key_type = reader.key_type();
27 let title = reader.header().title.clone();
28 let string_table = reader.read_string_table()?;
29 let frames = reader.read_all_frames()?;
30 Ok(Self {
31 schema,
32 title,
33 string_table,
34 frames,
35 key_type,
36 })
37 }
38
39 pub fn new(schema: Schema, title: impl Into<String>) -> Result<Self> {
40 let key_type = KeyType::from_field(schema.key_field())?;
41 Ok(Self {
42 schema,
43 title: title.into(),
44 string_table: Vec::new(),
45 frames: Vec::new(),
46 key_type,
47 })
48 }
49
50 pub fn schema(&self) -> &Schema {
51 &self.schema
52 }
53
54 pub fn title(&self) -> &str {
55 &self.title
56 }
57
58 pub fn frame_count(&self) -> u64 {
59 self.frames.len() as u64
60 }
61
62 pub fn string_table(&self) -> &[String] {
63 &self.string_table
64 }
65
66 pub fn frames(&self) -> &[OwnedFrame] {
67 &self.frames
68 }
69
70 pub fn append_string(&mut self, value: impl Into<String>) -> u32 {
71 let index = self.string_table.len() as u32;
72 self.string_table.push(value.into());
73 index
74 }
75
76 pub fn append_frame(&mut self, bytes: &[u8]) -> Result<()> {
77 let frame = OwnedFrame::new(&self.schema, bytes.to_vec())?;
78 let key = frame.as_ref().key(&self.schema, self.key_type)?;
79 if let Some(last) = self.last_key()? {
80 if key < last {
81 return Err(V1Error::KeyOrderViolation {
82 index: self.frames.len() as u64,
83 });
84 }
85 }
86 self.frames.push(frame);
87 Ok(())
88 }
89
90 pub fn append_frames<I, B>(&mut self, frames: I) -> Result<()>
91 where
92 I: IntoIterator<Item = B>,
93 B: AsRef<[u8]>,
94 {
95 for frame in frames {
96 self.append_frame(frame.as_ref())?;
97 }
98 Ok(())
99 }
100
101 pub fn delete_all_frames(&mut self) -> u64 {
102 let removed = self.frames.len() as u64;
103 self.frames.clear();
104 removed
105 }
106
107 pub fn delete_frames_before(&mut self, last_key: Key) -> Result<u64> {
108 let end = self.upper_bound(last_key)?;
109 self.frames.drain(0..end);
110 Ok(end as u64)
111 }
112
113 pub fn delete_frames_after(&mut self, first_key: Key) -> Result<u64> {
114 let begin = self.lower_bound(first_key)?;
115 let removed = self.frames.len() - begin;
116 self.frames.truncate(begin);
117 Ok(removed as u64)
118 }
119
120 pub fn delete_frames_between(&mut self, first_key: Key, last_key: Key) -> Result<u64> {
121 if first_key > last_key {
122 return Ok(0);
123 }
124 let begin = self.lower_bound(first_key)?;
125 let end = self.upper_bound(last_key)?;
126 let removed = end - begin;
127 self.frames.drain(begin..end);
128 Ok(removed as u64)
129 }
130
131 pub fn delete_frames<I>(&mut self, keys: I) -> Result<u64>
132 where
133 I: IntoIterator<Item = Key>,
134 {
135 let mut removed = 0u64;
136 for key in keys {
137 removed += self.delete_frames_between(key, key)?;
138 }
139 Ok(removed)
140 }
141
142 pub fn save_as(&self, path: impl AsRef<Path>) -> Result<()> {
143 let mut options = WriterOptions::new(self.title.clone());
144 let estimated_string_bytes: usize = self.string_table.iter().map(|s| s.len() + 5).sum();
145 options.string_table_preserved_length = estimated_string_bytes.max(1834) as u32;
146 let file = File::create(path)?;
147 let mut writer = Writer::new(file, self.schema.clone(), options)?;
148 for value in &self.string_table {
149 writer.append_string(value)?;
150 }
151 for frame in &self.frames {
152 writer.append_frame(frame.bytes())?;
153 }
154 Ok(())
155 }
156
157 fn last_key(&self) -> Result<Option<Key>> {
158 Ok(self
159 .frames
160 .last()
161 .map(|frame| frame.as_ref().key(&self.schema, self.key_type))
162 .transpose()?)
163 }
164
165 fn lower_bound(&self, key: Key) -> Result<usize> {
166 let mut lo = 0usize;
167 let mut hi = self.frames.len();
168 while lo < hi {
169 let mid = lo + ((hi - lo) >> 1);
170 let mid_key = self.frames[mid].as_ref().key(&self.schema, self.key_type)?;
171 if mid_key < key {
172 lo = mid + 1;
173 } else {
174 hi = mid;
175 }
176 }
177 Ok(lo)
178 }
179
180 fn upper_bound(&self, key: Key) -> Result<usize> {
181 let mut lo = 0usize;
182 let mut hi = self.frames.len();
183 while lo < hi {
184 let mid = lo + ((hi - lo) >> 1);
185 let mid_key = self.frames[mid].as_ref().key(&self.schema, self.key_type)?;
186 if mid_key <= key {
187 lo = mid + 1;
188 } else {
189 hi = mid;
190 }
191 }
192 Ok(lo)
193 }
194}