ion_rs/element/
element_stream_writer.rs

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// Represents a level into which the writer has stepped.
12// A writer that has not yet called step_in() is at the top level.
13#[derive(Debug, PartialEq)]
14struct ContainerContext {
15    annotations: Annotations,
16    field_name: Option<Symbol>,
17    container: Container,
18}
19
20// A container for values that the writer will write into.
21// A writer that has not yet called step_in() is at the top level.
22#[derive(Debug, PartialEq)]
23enum Container {
24    TopLevel,
25    SExpression(Vec<Element>),
26    List(Vec<Element>),
27    Struct(Vec<(Symbol, Element)>),
28}
29
30/// An application-level Element writer. This writer creates [`Element`]s, [`Extend::extend`]ing
31/// its output with each top-level value.
32pub 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` is never empty; it always has at least the top level.
63        self.containers.last().unwrap()
64    }
65
66    fn current_container_mut(&mut self) -> &mut ContainerContext {
67        // `self.containers` is never empty; it always has at least the top level.
68        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        // `self.containers` is never empty; it always has at least the top level.
91        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        // Nothing to do here
154        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                // TODO: This panic should be fixed as part of https://github.com/amazon-ion/ion-rust/issues/564
174                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                // TODO: This panic should be fixed as part of https://github.com/amazon-ion/ion-rust/issues/564
243                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()); // 100,000 arbitrarily chosen as 1/3 of the 15 decimal digits of precision
425            },
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(&timestamp),
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(&timestamp),
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(&timestamp),
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(&timestamp),
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(&timestamp),
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(&timestamp),
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}