1use crate::element::builders::StructBuilder;
2use crate::element::{Annotations, Element, IntoAnnotatedElement, Value};
3use crate::raw_symbol_token_ref::AsRawSymbolTokenRef;
4use crate::result::illegal_operation;
5use crate::types::Bytes;
6use crate::{
7 Decimal, Int, IonResult, IonType, IonWriter, RawSymbolTokenRef, Str, Symbol, SymbolTable,
8 Timestamp,
9};
10
11#[derive(Debug, PartialEq)]
14struct ContainerContext {
15 annotations: Annotations,
16 field_name: Option<Symbol>,
17 container: Container,
18}
19
20#[derive(Debug, PartialEq)]
23enum Container {
24 TopLevel,
25 SExpression(Vec<Element>),
26 List(Vec<Element>),
27 Struct(Vec<(Symbol, Element)>),
28}
29
30pub struct ElementStreamWriter<'e, E>
33where
34 E: Extend<Element>,
35{
36 output: &'e mut E,
37 containers: Vec<ContainerContext>,
38 annotations: Annotations,
39 field_name: Option<Symbol>,
40 symbol_table: SymbolTable,
41}
42
43impl<'e, E> ElementStreamWriter<'e, E>
44where
45 E: Extend<Element>,
46{
47 pub fn new(output: &'e mut E) -> ElementStreamWriter<'e, E> {
48 ElementStreamWriter {
49 output,
50 containers: vec![ContainerContext {
51 annotations: Annotations::empty(),
52 field_name: None,
53 container: Container::TopLevel,
54 }],
55 annotations: Annotations::empty(),
56 field_name: None,
57 symbol_table: SymbolTable::default(),
58 }
59 }
60
61 fn current_container(&self) -> &ContainerContext {
62 self.containers.last().unwrap()
64 }
65
66 fn current_container_mut(&mut self) -> &mut ContainerContext {
67 self.containers.last_mut().unwrap()
69 }
70
71 fn take_current_annotations(&mut self) -> Annotations {
72 std::mem::replace(&mut self.annotations, Annotations::empty())
73 }
74
75 fn push_container(&mut self, container: Container) -> IonResult<()> {
76 let annotations = self.take_current_annotations();
77 let field_name = self.field_name.take();
78 self.containers.push(ContainerContext {
79 annotations,
80 field_name,
81 container,
82 });
83 Ok(())
84 }
85
86 fn pop_container(&mut self) -> IonResult<ContainerContext> {
87 if self.containers.len() <= 1 {
88 return illegal_operation("cannot step out of the top level");
89 }
90 Ok(self.containers.pop().unwrap())
92 }
93
94 fn set_current_annotations(&mut self, annotations: Annotations) {
95 self.annotations = annotations;
96 }
97
98 fn write_scalar(&mut self, value: Value) -> IonResult<()> {
99 let annotations = self.take_current_annotations();
100 let element = value.with_annotations(annotations);
101 self.write_element(element)
102 }
103
104 fn write_element(&mut self, element: Element) -> IonResult<()> {
105 let field_name = self.field_name.take();
106 match &mut self.current_container_mut().container {
107 Container::TopLevel => self.output_mut().extend(std::iter::once(element)),
108 Container::SExpression(seq) => seq.push(element),
109 Container::List(seq) => seq.push(element),
110 Container::Struct(fields) => {
111 if let Some(field_name) = field_name {
112 fields.push((field_name, element));
113 } else {
114 return illegal_operation("Values inside a struct must have a field name.");
115 }
116 }
117 };
118 Ok(())
119 }
120
121 fn to_symbol<A>(&self, sym_ref: A) -> IonResult<Symbol>
122 where
123 A: AsRawSymbolTokenRef,
124 {
125 match sym_ref.as_raw_symbol_token_ref() {
126 RawSymbolTokenRef::SymbolId(symbol_id) => {
127 if self.symbol_table.sid_is_valid(symbol_id) {
128 Ok(self
129 .symbol_table
130 .symbol_for(symbol_id)
131 .cloned()
132 .unwrap_or(Symbol::unknown_text()))
133 } else {
134 illegal_operation(format!("Symbol ID ${symbol_id} is undefined."))
135 }
136 }
137 RawSymbolTokenRef::Text(txt) => Ok(Symbol::owned(txt)),
138 }
139 }
140}
141
142impl<'e, E> IonWriter for ElementStreamWriter<'e, E>
143where
144 E: Extend<Element>,
145{
146 type Output = &'e mut E;
147
148 fn ion_version(&self) -> (u8, u8) {
149 (1, 0)
150 }
151
152 fn write_ion_version_marker(&mut self, _major: u8, _minor: u8) -> IonResult<()> {
153 Ok(())
155 }
156
157 fn supports_text_symbol_tokens(&self) -> bool {
158 true
159 }
160
161 fn set_annotations<I, A>(&mut self, annotations: I)
162 where
163 A: AsRawSymbolTokenRef,
164 I: IntoIterator<Item = A>,
165 {
166 let annotations: IonResult<_> =
167 annotations.into_iter().map(|s| self.to_symbol(s)).collect();
168 match annotations {
169 Ok(annotations) => {
170 self.set_current_annotations(Annotations::new(annotations));
171 }
172 Err(e) => {
173 panic!("Cannot set as annotation due to {e}");
175 }
176 }
177 }
178
179 fn write_null(&mut self, ion_type: IonType) -> IonResult<()> {
180 self.write_scalar(Value::Null(ion_type))
181 }
182
183 fn write_bool(&mut self, value: bool) -> IonResult<()> {
184 self.write_scalar(Value::Bool(value))
185 }
186
187 fn write_i64(&mut self, value: i64) -> IonResult<()> {
188 self.write_scalar(Value::Int(Int::I64(value)))
189 }
190
191 fn write_int(&mut self, value: &Int) -> IonResult<()> {
192 self.write_scalar(Value::Int(value.clone()))
193 }
194
195 fn write_f32(&mut self, value: f32) -> IonResult<()> {
196 self.write_scalar(Value::Float(value as f64))
197 }
198
199 fn write_f64(&mut self, value: f64) -> IonResult<()> {
200 self.write_scalar(Value::Float(value))
201 }
202
203 fn write_decimal(&mut self, value: &Decimal) -> IonResult<()> {
204 self.write_scalar(Value::Decimal(value.clone()))
205 }
206
207 fn write_timestamp(&mut self, value: &Timestamp) -> IonResult<()> {
208 self.write_scalar(Value::Timestamp(value.clone()))
209 }
210
211 fn write_symbol<A: AsRawSymbolTokenRef>(&mut self, value: A) -> IonResult<()> {
212 let symbol = self.to_symbol(value)?;
213 self.write_scalar(Value::Symbol(symbol))
214 }
215
216 fn write_string<A: AsRef<str>>(&mut self, value: A) -> IonResult<()> {
217 self.write_scalar(Value::String(Str::from(value.as_ref())))
218 }
219
220 fn write_clob<A: AsRef<[u8]>>(&mut self, value: A) -> IonResult<()> {
221 self.write_scalar(Value::Clob(Bytes::from(value.as_ref())))
222 }
223
224 fn write_blob<A: AsRef<[u8]>>(&mut self, value: A) -> IonResult<()> {
225 self.write_scalar(Value::Blob(Bytes::from(value.as_ref())))
226 }
227
228 fn step_in(&mut self, ion_type: IonType) -> IonResult<()> {
229 let container = match ion_type {
230 IonType::Struct => Container::Struct(vec![]),
231 IonType::List => Container::List(vec![]),
232 IonType::SExp => Container::SExpression(vec![]),
233 _ => return illegal_operation(format!("Cannot step into a(n) {ion_type:?}")),
234 };
235 self.push_container(container)
236 }
237
238 fn set_field_name<A: AsRawSymbolTokenRef>(&mut self, name: A) {
239 match self.to_symbol(name) {
240 Ok(s) => self.field_name = Some(s),
241 Err(e) => {
242 panic!("Cannot set field name due to {e}");
244 }
245 }
246 }
247
248 fn parent_type(&self) -> Option<IonType> {
249 match self.current_container().container {
250 Container::TopLevel => None,
251 Container::SExpression(_) => Some(IonType::SExp),
252 Container::List(_) => Some(IonType::List),
253 Container::Struct(_) => Some(IonType::Struct),
254 }
255 }
256
257 fn depth(&self) -> usize {
258 self.containers.len() - 1
259 }
260
261 fn step_out(&mut self) -> IonResult<()> {
262 let ContainerContext {
263 annotations,
264 field_name,
265 container,
266 } = self.pop_container()?;
267 let value = match container {
268 Container::TopLevel => return illegal_operation("cannot step out of the top level"),
269 Container::SExpression(seq) => Value::SExp(seq.into()),
270 Container::List(seq) => Value::List(seq.into()),
271 Container::Struct(fields) => {
272 Value::Struct(StructBuilder::new().with_fields(fields).build())
273 }
274 };
275 self.field_name = field_name;
276 let element = value.with_annotations(annotations);
277 self.write_element(element)
278 }
279
280 fn flush(&mut self) -> IonResult<()> {
281 Ok(())
282 }
283
284 fn output(&self) -> &Self::Output {
285 &self.output
286 }
287
288 fn output_mut(&mut self) -> &mut Self::Output {
289 &mut self.output
290 }
291}
292
293#[cfg(test)]
294mod tests {
295
296 use crate::element::element_stream_writer::ElementStreamWriter;
297 use crate::element::{Element, IntoAnnotatedElement, Value};
298
299 use crate::element::builders::{SequenceBuilder, StructBuilder};
300 use crate::result::IonResult;
301
302 use crate::types::{Bytes, Timestamp};
303 use crate::writer::IonWriter;
304 use crate::{Decimal, Int, IonType, Symbol};
305
306 #[track_caller]
307 fn writer_test_with_assertion<F, E, A>(mut commands: F, expected: Vec<E>, mut assertion: A)
308 where
309 E: Into<Element>,
310 F: FnMut(&mut ElementStreamWriter<Vec<Element>>) -> IonResult<()>,
311 A: FnMut(Vec<Element>, Vec<Element>),
312 {
313 let mut output = Vec::new();
314 let mut writer = ElementStreamWriter::new(&mut output);
315 commands(&mut writer).expect("Invalid ElementStreamWriter test commands.");
316 writer.flush().expect("Call to flush() failed.");
317 let expected: Vec<Element> = expected.into_iter().map(|e| e.into()).collect();
318
319 assertion(output, expected);
320 }
321
322 #[track_caller]
323 fn writer_test<F, E>(commands: F, expected: Vec<E>)
324 where
325 E: Into<Element>,
326 F: FnMut(&mut ElementStreamWriter<Vec<Element>>) -> IonResult<()>,
327 {
328 writer_test_with_assertion(commands, expected, |actual, expected| {
329 assert_eq!(actual, expected);
330 })
331 }
332
333 #[track_caller]
334 fn writer_scalar_test_with_assertion<F, E, A>(commands: F, expected: E, mut assertion: A)
335 where
336 E: Into<Element>,
337 F: FnMut(&mut ElementStreamWriter<Vec<Element>>) -> IonResult<()>,
338 A: FnMut(Element, Element),
339 {
340 writer_test_with_assertion(
341 commands,
342 vec![expected.into()],
343 |mut actual, mut expected| {
344 assert_eq!(actual.len(), 1);
345 assert_eq!(expected.len(), 1);
346 assertion(actual.pop().unwrap(), expected.pop().unwrap());
347 },
348 );
349 }
350
351 #[track_caller]
352 fn writer_scalar_test<F, E>(commands: F, expected: E)
353 where
354 E: Into<Element>,
355 F: FnMut(&mut ElementStreamWriter<Vec<Element>>) -> IonResult<()>,
356 {
357 writer_scalar_test_with_assertion(commands, expected, |actual, expected| {
358 assert_eq!(actual, expected);
359 })
360 }
361
362 #[test]
363 fn write_null_null() {
364 writer_scalar_test(|w| w.write_null(IonType::Null), Value::Null(IonType::Null));
365 }
366
367 #[test]
368 fn write_null_string() {
369 writer_scalar_test(
370 |w| w.write_null(IonType::String),
371 Value::Null(IonType::String),
372 );
373 }
374
375 #[test]
376 fn write_bool_true() {
377 writer_scalar_test(|w| w.write_bool(true), Value::Bool(true));
378 }
379
380 #[test]
381 fn write_bool_false() {
382 writer_scalar_test(|w| w.write_bool(false), Value::Bool(false));
383 }
384
385 #[test]
386 fn write_i64() {
387 writer_scalar_test(|w| w.write_i64(7), Value::Int(Int::I64(7)));
388 }
389
390 #[test]
391 fn write_f32() {
392 writer_scalar_test(|w| w.write_f32(700f32), Value::Float(700.));
393 }
394
395 #[test]
396 fn write_f64() {
397 writer_scalar_test(|w| w.write_f64(700f64), Value::Float(700.));
398 }
399
400 #[test]
401 fn write_annotated_i64() {
402 writer_scalar_test(
403 |w| {
404 w.set_annotations(["foo", "bar", "baz quux"]);
405 w.write_i64(7)
406 },
407 Value::Int(Int::I64(7)).with_annotations(["foo", "bar", "baz quux"]),
408 );
409 }
410
411 #[test]
412 fn write_decimal() {
413 let decimal: Decimal = 731221.9948f64.try_into().expect("float to decimal");
414 writer_scalar_test_with_assertion(
415 |w| w.write_decimal(&decimal),
416 Value::Decimal(Decimal::new(7312219948u64, -4)),
417 |actual, expected| {
418 assert_eq!(actual.ion_type(), IonType::Decimal);
419 assert_eq!(expected.ion_type(), IonType::Decimal);
420 let actual = actual.as_decimal().unwrap().clone();
421 let expected = expected.as_decimal().unwrap().clone();
422 let (diff_int, diff_fract) = Decimal::difference_by_parts_lossy(&actual, &expected);
423 assert_eq!(diff_int, 0.into(), "integer component expected equal");
424 assert!(diff_fract < 100_000.into()); },
426 );
427 }
428
429 #[test]
430 fn write_timestamp_with_year() {
431 let timestamp = Timestamp::with_year(2000)
432 .build()
433 .expect("building timestamp failed");
434 writer_scalar_test(
435 |w| w.write_timestamp(×tamp),
436 Value::Timestamp(
437 Timestamp::with_year(2000)
438 .build()
439 .expect("timestamp expected value"),
440 ),
441 );
442 }
443
444 #[test]
445 fn write_timestamp_with_month() {
446 let timestamp = Timestamp::with_year(2000)
447 .with_month(8)
448 .build()
449 .expect("building timestamp failed");
450 writer_scalar_test(
451 |w| w.write_timestamp(×tamp),
452 Value::Timestamp(
453 Timestamp::with_year(2000)
454 .with_month(8)
455 .build()
456 .expect("timestamp expected value"),
457 ),
458 );
459 }
460
461 #[test]
462 fn write_timestamp_with_ymd() {
463 let timestamp = Timestamp::with_ymd(2000, 8, 22)
464 .build()
465 .expect("building timestamp failed");
466 writer_scalar_test(
467 |w| w.write_timestamp(×tamp),
468 Value::Timestamp(
469 Timestamp::with_year(2000)
470 .with_month(8)
471 .with_day(22)
472 .build()
473 .expect("timestamp expected value"),
474 ),
475 );
476 }
477
478 #[test]
479 fn write_timestamp_with_ymd_hms() {
480 let timestamp = Timestamp::with_ymd(2000, 8, 22)
481 .with_hms(15, 45, 11)
482 .build_at_offset(2 * 60)
483 .expect("building timestamp failed");
484 writer_scalar_test(
485 |w| w.write_timestamp(×tamp),
486 Value::Timestamp(
487 Timestamp::with_year(2000)
488 .with_month(8)
489 .with_day(22)
490 .with_hms(15, 45, 11)
491 .build_at_offset(120)
492 .expect("timestamp expected value"),
493 ),
494 );
495 }
496
497 #[test]
498 fn write_timestamp_with_ymd_hms_millis() {
499 let timestamp = Timestamp::with_ymd_hms_millis(2000, 8, 22, 15, 45, 11, 931)
500 .build_at_offset(-5 * 60)
501 .expect("building timestamp failed");
502 writer_scalar_test(
503 |w| w.write_timestamp(×tamp),
504 Value::Timestamp(
505 Timestamp::with_year(2000)
506 .with_month(8)
507 .with_day(22)
508 .with_hms(15, 45, 11)
509 .with_milliseconds(931)
510 .build_at_offset(-300)
511 .expect("timestamp expected value"),
512 ),
513 );
514 }
515
516 #[test]
517 fn write_timestamp_with_ymd_hms_millis_unknown_offset() {
518 let timestamp = Timestamp::with_ymd_hms_millis(2000, 8, 22, 15, 45, 11, 931)
519 .build_at_unknown_offset()
520 .expect("building timestamp failed");
521 writer_scalar_test(
522 |w| w.write_timestamp(×tamp),
523 Value::Timestamp(
524 Timestamp::with_year(2000)
525 .with_month(8)
526 .with_day(22)
527 .with_hms(15, 45, 11)
528 .with_milliseconds(931)
529 .build_at_unknown_offset()
530 .expect("timestamp expected value"),
531 ),
532 );
533 }
534
535 #[test]
536 fn write_blob() {
537 writer_scalar_test(
538 |w| w.write_blob("hello".as_bytes()),
539 Value::Blob(Bytes::from(&[104, 101, 108, 108, 111])),
540 );
541 }
542
543 #[test]
544 fn write_clob() {
545 writer_scalar_test(
546 |w| w.write_clob("a\"\'\n".as_bytes()),
547 Value::Clob(Bytes::from(&[97, 34, 39, 10])),
548 );
549 writer_scalar_test(
550 |w| w.write_clob("❤️".as_bytes()),
551 Value::Clob(Bytes::from("❤️")),
552 );
553 writer_scalar_test(
554 |w| w.write_clob("hello world".as_bytes()),
555 Value::Clob(Bytes::from(&[
556 104, 101, 108, 108, 111, 32, 119, 111, 114, 108, 100,
557 ])),
558 );
559 }
560
561 #[test]
562 fn top_level_values() {
563 writer_test(
564 |w| {
565 w.write_bool(true)?;
566 w.write_bool(false)
567 },
568 vec![Value::Bool(true), Value::Bool(false)],
569 );
570 }
571
572 #[test]
573 fn top_level_values_with_nesting() {
574 writer_test(
575 |w| {
576 w.write_bool(true)?;
577 w.step_in(IonType::List)?;
578 w.write_string("foo")?;
579 w.write_i64(21)?;
580 w.write_symbol("bar")?;
581 w.step_out()
582 },
583 vec![
584 Value::Bool(true),
585 Value::List(
586 SequenceBuilder::new()
587 .push("foo")
588 .push(21)
589 .push(Value::Symbol(Symbol::from("bar")))
590 .build(),
591 ),
592 ],
593 );
594 }
595
596 #[test]
597 fn write_list() {
598 writer_scalar_test(
599 |w| {
600 w.step_in(IonType::List)?;
601 w.write_string("foo")?;
602 w.write_i64(21)?;
603 w.write_symbol("bar")?;
604 w.step_out()
605 },
606 Value::List(
607 SequenceBuilder::new()
608 .push("foo")
609 .push(21)
610 .push(Value::Symbol(Symbol::from("bar")))
611 .build(),
612 ),
613 );
614 }
615
616 #[test]
617 fn write_nested_list() {
618 writer_scalar_test(
619 |w| {
620 w.step_in(IonType::List)?;
621 w.write_string("foo")?;
622 w.write_i64(21)?;
623 w.step_in(IonType::List)?;
624 w.write_symbol("bar")?;
625 w.step_out()?;
626 w.step_out()
627 },
628 Value::List(
629 SequenceBuilder::new()
630 .push("foo")
631 .push(21)
632 .push(Value::List(
633 SequenceBuilder::new()
634 .push(Value::Symbol(Symbol::from("bar")))
635 .build(),
636 ))
637 .build(),
638 ),
639 );
640 }
641
642 #[test]
643 fn write_s_expression() {
644 writer_scalar_test(
645 |w| {
646 w.step_in(IonType::SExp)?;
647 w.write_string("foo")?;
648 w.write_i64(21)?;
649 w.write_symbol("bar")?;
650 w.step_out()
651 },
652 Value::SExp(
653 SequenceBuilder::new()
654 .push("foo")
655 .push(21)
656 .push(Value::Symbol(Symbol::from("bar")))
657 .build(),
658 ),
659 );
660 }
661
662 #[test]
663 fn write_struct() {
664 writer_scalar_test(
665 |w| {
666 w.step_in(IonType::Struct)?;
667 w.set_field_name("a");
668 w.write_string("foo")?;
669 w.set_field_name("b");
670 w.write_i64(21)?;
671 w.set_field_name("c");
672 w.set_annotations(["quux"]);
673 w.write_symbol("bar")?;
674 w.step_out()
675 },
676 Value::Struct(
677 StructBuilder::new()
678 .with_field("a", "foo")
679 .with_field("b", 21)
680 .with_field(
681 "c",
682 Value::Symbol(Symbol::from("bar")).with_annotations(["quux"]),
683 )
684 .build(),
685 ),
686 );
687 }
688}