Skip to main content

grib_reader/
lib.rs

1//! Pure-Rust GRIB file reader.
2//!
3//! The current implementation supports the production-critical baseline for both
4//! GRIB1 and GRIB2: regular latitude/longitude grids, simple packing, and
5//! GRIB2 complex packing with general group splitting.
6//!
7//! # Example
8//!
9//! ```no_run
10//! use grib_reader::GribFile;
11//!
12//! let file = GribFile::open("gfs.grib2")?;
13//! println!("messages: {}", file.message_count());
14//!
15//! for msg in file.messages() {
16//!     println!(
17//!         "  {} {:?} {:04}-{:02}-{:02}T{:02}:{:02}:{:02}Z",
18//!         msg.parameter_name(),
19//!         msg.grid_shape(),
20//!         msg.reference_time().year,
21//!         msg.reference_time().month,
22//!         msg.reference_time().day,
23//!         msg.reference_time().hour,
24//!         msg.reference_time().minute,
25//!         msg.reference_time().second,
26//!     );
27//! }
28//!
29//! let data = file.message(0)?.read_data_as_f64()?;
30//! println!("shape: {:?}", data.shape());
31//! # Ok::<(), grib_reader::Error>(())
32//! ```
33
34pub mod data;
35pub mod error;
36pub mod grib1;
37pub mod grid;
38pub mod indicator;
39pub mod metadata;
40pub mod parameter;
41pub mod product;
42pub mod sections;
43mod util;
44
45pub use data::DecodeSample;
46pub use error::{Error, Result};
47pub use metadata::{ForecastTimeUnit, Parameter, ReferenceTime};
48pub use product::{
49    AnalysisOrForecastTemplate, FixedSurface, Identification, ProductDefinition,
50    ProductDefinitionTemplate,
51};
52
53use std::path::Path;
54
55use memmap2::Mmap;
56use ndarray::{ArrayD, IxDyn};
57
58use crate::data::{
59    bitmap_payload as grib2_bitmap_payload, count_bitmap_present_points, decode_field_into,
60    decode_payload_into, DataRepresentation,
61};
62use crate::grib1::{BinaryDataSection, GridDescription};
63use crate::grid::GridDefinition;
64use crate::indicator::Indicator;
65use crate::sections::{index_fields, FieldSections, SectionRef};
66
67#[cfg(feature = "rayon")]
68use rayon::prelude::*;
69
70const GRIB_MAGIC: &[u8; 4] = b"GRIB";
71
72/// Configuration for opening GRIB data.
73#[derive(Debug, Clone, Copy)]
74pub struct OpenOptions {
75    /// When `true`, the first malformed GRIB candidate aborts opening.
76    ///
77    /// When `false`, candidate offsets with invalid framing are skipped and
78    /// scanning continues. Once a candidate has a valid indicator, message
79    /// length, and end marker, any indexing or decoding error is returned.
80    pub strict: bool,
81}
82
83impl Default for OpenOptions {
84    fn default() -> Self {
85        Self { strict: true }
86    }
87}
88
89/// A parsed GRIB field.
90#[derive(Debug, Clone)]
91pub struct MessageMetadata {
92    pub edition: u8,
93    pub center_id: u16,
94    pub subcenter_id: u16,
95    pub discipline: Option<u8>,
96    pub reference_time: ReferenceTime,
97    pub parameter: Parameter,
98    pub grid: GridDefinition,
99    pub data_representation: DataRepresentation,
100    pub forecast_time_unit: Option<u8>,
101    pub forecast_time: Option<u32>,
102    pub message_offset: u64,
103    pub message_length: u64,
104    pub field_index_in_message: usize,
105    grib1_product: Option<grib1::ProductDefinition>,
106    grib2_identification: Option<Identification>,
107    grib2_product: Option<ProductDefinition>,
108}
109
110/// A GRIB file containing one or more logical fields.
111pub struct GribFile {
112    data: GribData,
113    messages: Vec<MessageIndex>,
114}
115
116#[derive(Clone)]
117struct MessageIndex {
118    offset: usize,
119    length: usize,
120    metadata: MessageMetadata,
121    decode_plan: DecodePlan,
122}
123
124#[derive(Clone, Copy)]
125enum DecodePlan {
126    Grib1 {
127        bitmap: Option<SectionRef>,
128        data: SectionRef,
129    },
130    Grib2(FieldSections),
131}
132
133enum GribData {
134    Mmap(Mmap),
135    Bytes(Vec<u8>),
136}
137
138impl GribData {
139    fn as_bytes(&self) -> &[u8] {
140        match self {
141            GribData::Mmap(m) => m,
142            GribData::Bytes(b) => b,
143        }
144    }
145}
146
147impl GribFile {
148    /// Open a GRIB file from disk using memory-mapped I/O.
149    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
150        Self::open_with_options(path, OpenOptions::default())
151    }
152
153    /// Open a GRIB file from disk using explicit decoder options.
154    pub fn open_with_options<P: AsRef<Path>>(path: P, options: OpenOptions) -> Result<Self> {
155        let file = std::fs::File::open(path.as_ref())
156            .map_err(|e| Error::Io(e, path.as_ref().display().to_string()))?;
157        let mmap = unsafe { Mmap::map(&file) }
158            .map_err(|e| Error::Io(e, path.as_ref().display().to_string()))?;
159        Self::from_data(GribData::Mmap(mmap), options)
160    }
161
162    /// Open a GRIB file from an owned byte buffer.
163    pub fn from_bytes(data: Vec<u8>) -> Result<Self> {
164        Self::from_bytes_with_options(data, OpenOptions::default())
165    }
166
167    /// Open a GRIB file from an owned byte buffer using explicit decoder options.
168    pub fn from_bytes_with_options(data: Vec<u8>, options: OpenOptions) -> Result<Self> {
169        Self::from_data(GribData::Bytes(data), options)
170    }
171
172    fn from_data(data: GribData, options: OpenOptions) -> Result<Self> {
173        let messages = scan_messages(data.as_bytes(), options)?;
174        if messages.is_empty() {
175            return Err(Error::NoMessages);
176        }
177        Ok(Self { data, messages })
178    }
179
180    /// Returns the GRIB edition of the first field.
181    pub fn edition(&self) -> u8 {
182        self.messages
183            .first()
184            .map(|message| message.metadata.edition)
185            .unwrap_or(0)
186    }
187
188    /// Returns the number of logical fields in the file.
189    pub fn message_count(&self) -> usize {
190        self.messages.len()
191    }
192
193    /// Access a field by index.
194    pub fn message(&self, index: usize) -> Result<Message<'_>> {
195        let record = self
196            .messages
197            .get(index)
198            .ok_or(Error::MessageNotFound(index))?;
199        let bytes = &self.data.as_bytes()[record.offset..record.offset + record.length];
200        Ok(Message {
201            index,
202            bytes,
203            metadata: &record.metadata,
204            decode_plan: record.decode_plan,
205        })
206    }
207
208    /// Iterate over all fields.
209    pub fn messages(&self) -> MessageIter<'_> {
210        MessageIter {
211            file: self,
212            index: 0,
213        }
214    }
215
216    /// Decode every field in the file.
217    pub fn read_all_data_as_f64(&self) -> Result<Vec<ArrayD<f64>>> {
218        #[cfg(feature = "rayon")]
219        {
220            (0..self.message_count())
221                .into_par_iter()
222                .map(|index| self.message(index)?.read_data_as_f64())
223                .collect()
224        }
225        #[cfg(not(feature = "rayon"))]
226        {
227            (0..self.message_count())
228                .map(|index| self.message(index)?.read_data_as_f64())
229                .collect()
230        }
231    }
232
233    /// Decode every field in the file as `f32`.
234    pub fn read_all_data_as_f32(&self) -> Result<Vec<ArrayD<f32>>> {
235        #[cfg(feature = "rayon")]
236        {
237            (0..self.message_count())
238                .into_par_iter()
239                .map(|index| self.message(index)?.read_data_as_f32())
240                .collect()
241        }
242        #[cfg(not(feature = "rayon"))]
243        {
244            (0..self.message_count())
245                .map(|index| self.message(index)?.read_data_as_f32())
246                .collect()
247        }
248    }
249}
250
251/// A single logical GRIB field.
252pub struct Message<'a> {
253    bytes: &'a [u8],
254    metadata: &'a MessageMetadata,
255    decode_plan: DecodePlan,
256    index: usize,
257}
258
259impl<'a> Message<'a> {
260    pub fn edition(&self) -> u8 {
261        self.metadata.edition
262    }
263
264    pub fn index(&self) -> usize {
265        self.index
266    }
267
268    pub fn metadata(&self) -> &MessageMetadata {
269        self.metadata
270    }
271
272    pub fn reference_time(&self) -> &ReferenceTime {
273        &self.metadata.reference_time
274    }
275
276    pub fn parameter(&self) -> &Parameter {
277        &self.metadata.parameter
278    }
279
280    pub fn center_id(&self) -> u16 {
281        self.metadata.center_id
282    }
283
284    pub fn subcenter_id(&self) -> u16 {
285        self.metadata.subcenter_id
286    }
287
288    pub fn identification(&self) -> Option<&Identification> {
289        self.metadata.grib2_identification.as_ref()
290    }
291
292    pub fn product_definition(&self) -> Option<&ProductDefinition> {
293        self.metadata.grib2_product.as_ref()
294    }
295
296    pub fn grib1_product_definition(&self) -> Option<&grib1::ProductDefinition> {
297        self.metadata.grib1_product.as_ref()
298    }
299
300    pub fn grid_definition(&self) -> &GridDefinition {
301        &self.metadata.grid
302    }
303
304    pub fn parameter_name(&self) -> &'static str {
305        self.metadata.parameter.short_name
306    }
307
308    pub fn parameter_description(&self) -> &'static str {
309        self.metadata.parameter.description
310    }
311
312    pub fn forecast_time_unit(&self) -> Option<u8> {
313        self.metadata.forecast_time_unit
314    }
315
316    pub fn forecast_time_unit_kind(&self) -> Option<ForecastTimeUnit> {
317        let unit = self.metadata.forecast_time_unit?;
318        ForecastTimeUnit::from_edition_and_code(self.metadata.edition, unit)
319    }
320
321    pub fn forecast_time(&self) -> Option<u32> {
322        self.metadata.forecast_time
323    }
324
325    pub fn valid_time(&self) -> Option<ReferenceTime> {
326        let unit = self.metadata.forecast_time_unit?;
327        let lead = self.metadata.forecast_time?;
328        self.metadata
329            .reference_time
330            .checked_add_forecast_time_by_edition(self.metadata.edition, unit, lead)
331    }
332
333    pub fn grid_shape(&self) -> (usize, usize) {
334        self.metadata.grid.shape()
335    }
336
337    pub fn latitudes(&self) -> Option<Vec<f64>> {
338        match &self.metadata.grid {
339            GridDefinition::LatLon(grid) => Some(grid.latitudes()),
340            GridDefinition::Unsupported(_) => None,
341        }
342    }
343
344    pub fn longitudes(&self) -> Option<Vec<f64>> {
345        match &self.metadata.grid {
346            GridDefinition::LatLon(grid) => Some(grid.longitudes()),
347            GridDefinition::Unsupported(_) => None,
348        }
349    }
350
351    pub fn decode_into<T: DecodeSample>(&self, out: &mut [T]) -> Result<()> {
352        let grid = match &self.metadata.grid {
353            GridDefinition::LatLon(grid) => grid,
354            GridDefinition::Unsupported(template) => {
355                return Err(Error::UnsupportedGridTemplate(*template));
356            }
357        };
358
359        match self.decode_plan {
360            DecodePlan::Grib2(field) => {
361                let data_section = section_bytes(self.bytes, field.data);
362                let bitmap_section = match field.bitmap {
363                    Some(section) => grib2_bitmap_payload(section_bytes(self.bytes, section))?,
364                    None => None,
365                };
366                decode_field_into(
367                    data_section,
368                    &self.metadata.data_representation,
369                    bitmap_section,
370                    self.metadata.grid.num_points(),
371                    out,
372                )?
373            }
374            DecodePlan::Grib1 { bitmap, data } => {
375                let bitmap_section = match bitmap {
376                    Some(section) => grib1::bitmap_payload(section_bytes(self.bytes, section))?,
377                    None => None,
378                };
379                let data_section = section_bytes(self.bytes, data);
380                if data_section.len() < 11 {
381                    return Err(Error::InvalidSection {
382                        section: 4,
383                        reason: format!("expected at least 11 bytes, got {}", data_section.len()),
384                    });
385                }
386                decode_payload_into(
387                    &data_section[11..],
388                    &self.metadata.data_representation,
389                    bitmap_section,
390                    self.metadata.grid.num_points(),
391                    out,
392                )?
393            }
394        }
395
396        grid.reorder_for_ndarray_in_place(out)
397    }
398
399    pub fn read_flat_data_as_f64(&self) -> Result<Vec<f64>> {
400        let mut decoded = vec![0.0; self.metadata.grid.num_points()];
401        self.decode_into(&mut decoded)?;
402        Ok(decoded)
403    }
404
405    pub fn read_flat_data_as_f32(&self) -> Result<Vec<f32>> {
406        let mut decoded = vec![0.0_f32; self.metadata.grid.num_points()];
407        self.decode_into(&mut decoded)?;
408        Ok(decoded)
409    }
410
411    pub fn read_data_as_f64(&self) -> Result<ArrayD<f64>> {
412        let ordered = self.read_flat_data_as_f64()?;
413        ArrayD::from_shape_vec(IxDyn(&self.metadata.grid.ndarray_shape()), ordered)
414            .map_err(|e| Error::Other(format!("failed to build ndarray from decoded field: {e}")))
415    }
416
417    pub fn read_data_as_f32(&self) -> Result<ArrayD<f32>> {
418        let ordered = self.read_flat_data_as_f32()?;
419        ArrayD::from_shape_vec(IxDyn(&self.metadata.grid.ndarray_shape()), ordered)
420            .map_err(|e| Error::Other(format!("failed to build ndarray from decoded field: {e}")))
421    }
422
423    pub fn raw_bytes(&self) -> &[u8] {
424        self.bytes
425    }
426}
427
428/// Iterator over fields in a GRIB file.
429pub struct MessageIter<'a> {
430    file: &'a GribFile,
431    index: usize,
432}
433
434impl<'a> Iterator for MessageIter<'a> {
435    type Item = Message<'a>;
436
437    fn next(&mut self) -> Option<Self::Item> {
438        if self.index >= self.file.message_count() {
439            return None;
440        }
441        let message = self.file.message(self.index).ok()?;
442        self.index += 1;
443        Some(message)
444    }
445}
446
447fn section_bytes(msg_bytes: &[u8], section: SectionRef) -> &[u8] {
448    &msg_bytes[section.offset..section.offset + section.length]
449}
450
451fn scan_messages(data: &[u8], options: OpenOptions) -> Result<Vec<MessageIndex>> {
452    let mut messages = Vec::new();
453    let mut pos = 0usize;
454
455    while pos + 8 <= data.len() {
456        if &data[pos..pos + 4] != GRIB_MAGIC {
457            pos += 1;
458            continue;
459        }
460
461        let (indicator, next_pos) = match locate_message(data, pos) {
462            Ok(located) => located,
463            Err(err) if !options.strict && is_recoverable_candidate_error(&err) => {
464                pos += 4;
465                continue;
466            }
467            Err(err) => return Err(err),
468        };
469
470        let message_bytes = &data[pos..next_pos];
471        let indexed = match indicator.edition {
472            1 => index_grib1_message(message_bytes, pos)?,
473            2 => index_grib2_message(message_bytes, pos, &indicator)?,
474            other => return Err(Error::UnsupportedEdition(other)),
475        };
476
477        messages.extend(indexed);
478        pos = next_pos;
479    }
480
481    Ok(messages)
482}
483
484fn is_recoverable_candidate_error(err: &Error) -> bool {
485    matches!(err, Error::InvalidMessage(_) | Error::Truncated { .. })
486}
487
488fn locate_message(data: &[u8], pos: usize) -> Result<(Indicator, usize)> {
489    let indicator = Indicator::parse(&data[pos..]).ok_or_else(|| {
490        Error::InvalidMessage(format!("failed to parse indicator at byte offset {pos}"))
491    })?;
492    let length = indicator.total_length as usize;
493    if length < 12 {
494        return Err(Error::InvalidMessage(format!(
495            "message at byte offset {pos} reports impossible length {length}"
496        )));
497    }
498    let end = pos
499        .checked_add(length)
500        .ok_or_else(|| Error::InvalidMessage("message length overflow".into()))?;
501    if end > data.len() {
502        return Err(Error::Truncated { offset: end as u64 });
503    }
504    if &data[end - 4..end] != b"7777" {
505        return Err(Error::InvalidMessage(format!(
506            "message at byte offset {pos} does not end with 7777"
507        )));
508    }
509
510    Ok((indicator, end))
511}
512
513fn index_grib1_message(message_bytes: &[u8], offset: usize) -> Result<Vec<MessageIndex>> {
514    let sections = grib1::parse_message_sections(message_bytes)?;
515    let grid_ref = sections.grid.ok_or_else(|| {
516        Error::InvalidSectionOrder(
517            "GRIB1 decoding requires an explicit grid definition section".into(),
518        )
519    })?;
520    let grid_description = GridDescription::parse(section_bytes(message_bytes, grid_ref))?;
521    let grid = grid_description.grid;
522
523    let bitmap_present_count = match sections.bitmap {
524        Some(bitmap) => {
525            let bitmap_payload = grib1::bitmap_payload(section_bytes(message_bytes, bitmap))?
526                .ok_or(Error::MissingBitmap)?;
527            count_bitmap_present_points(bitmap_payload, grid.num_points())?
528        }
529        None => grid.num_points(),
530    };
531
532    let (_bds, data_representation) = BinaryDataSection::parse(
533        section_bytes(message_bytes, sections.data),
534        sections.product.decimal_scale,
535        bitmap_present_count,
536    )?;
537    let parameter = sections.product.parameter();
538
539    Ok(vec![MessageIndex {
540        offset,
541        length: message_bytes.len(),
542        decode_plan: DecodePlan::Grib1 {
543            bitmap: sections.bitmap,
544            data: sections.data,
545        },
546        metadata: MessageMetadata {
547            edition: 1,
548            center_id: sections.product.center_id as u16,
549            subcenter_id: sections.product.subcenter_id as u16,
550            discipline: None,
551            reference_time: sections.product.reference_time,
552            parameter,
553            grid,
554            data_representation,
555            forecast_time_unit: Some(sections.product.forecast_time_unit),
556            forecast_time: sections.product.forecast_time(),
557            message_offset: offset as u64,
558            message_length: message_bytes.len() as u64,
559            field_index_in_message: 0,
560            grib1_product: Some(sections.product),
561            grib2_identification: None,
562            grib2_product: None,
563        },
564    }])
565}
566
567fn index_grib2_message(
568    message_bytes: &[u8],
569    offset: usize,
570    indicator: &Indicator,
571) -> Result<Vec<MessageIndex>> {
572    let fields = index_fields(message_bytes)?;
573    let mut messages = Vec::with_capacity(fields.len());
574
575    for (field_index_in_message, field_sections) in fields.into_iter().enumerate() {
576        let identification =
577            Identification::parse(section_bytes(message_bytes, field_sections.identification))?;
578        let grid = GridDefinition::parse(section_bytes(message_bytes, field_sections.grid))?;
579        let product =
580            ProductDefinition::parse(section_bytes(message_bytes, field_sections.product))?;
581        let data_representation = DataRepresentation::parse(section_bytes(
582            message_bytes,
583            field_sections.data_representation,
584        ))?;
585        let parameter = Parameter::new_grib2(
586            indicator.discipline,
587            product.parameter_category,
588            product.parameter_number,
589            product.parameter_name(indicator.discipline),
590            product.parameter_description(indicator.discipline),
591        );
592
593        messages.push(MessageIndex {
594            offset,
595            length: message_bytes.len(),
596            decode_plan: DecodePlan::Grib2(field_sections),
597            metadata: MessageMetadata {
598                edition: indicator.edition,
599                center_id: identification.center_id,
600                subcenter_id: identification.subcenter_id,
601                discipline: Some(indicator.discipline),
602                reference_time: ReferenceTime {
603                    year: identification.reference_year,
604                    month: identification.reference_month,
605                    day: identification.reference_day,
606                    hour: identification.reference_hour,
607                    minute: identification.reference_minute,
608                    second: identification.reference_second,
609                },
610                parameter,
611                grid,
612                data_representation,
613                forecast_time_unit: product.forecast_time_unit(),
614                forecast_time: product.forecast_time(),
615                message_offset: offset as u64,
616                message_length: message_bytes.len() as u64,
617                field_index_in_message,
618                grib1_product: None,
619                grib2_identification: Some(identification),
620                grib2_product: Some(product),
621            },
622        });
623    }
624
625    Ok(messages)
626}
627
628#[cfg(test)]
629mod tests {
630    use super::*;
631
632    fn grib_i32_bytes(value: i32) -> [u8; 4] {
633        if value >= 0 {
634            (value as u32).to_be_bytes()
635        } else {
636            ((-value) as u32 | 0x8000_0000).to_be_bytes()
637        }
638    }
639
640    fn build_indicator(total_len: usize, discipline: u8) -> Vec<u8> {
641        let mut indicator = Vec::with_capacity(16);
642        indicator.extend_from_slice(b"GRIB");
643        indicator.extend_from_slice(&[0, 0]);
644        indicator.push(discipline);
645        indicator.push(2);
646        indicator.extend_from_slice(&(total_len as u64).to_be_bytes());
647        indicator
648    }
649
650    fn build_identification() -> Vec<u8> {
651        let mut section = vec![0u8; 21];
652        section[..4].copy_from_slice(&(21u32).to_be_bytes());
653        section[4] = 1;
654        section[5..7].copy_from_slice(&7u16.to_be_bytes());
655        section[7..9].copy_from_slice(&0u16.to_be_bytes());
656        section[9] = 35;
657        section[10] = 1;
658        section[11] = 1;
659        section[12..14].copy_from_slice(&2026u16.to_be_bytes());
660        section[14] = 3;
661        section[15] = 20;
662        section[16] = 12;
663        section[17] = 0;
664        section[18] = 0;
665        section[19] = 0;
666        section[20] = 1;
667        section
668    }
669
670    fn build_grid(ni: u32, nj: u32, scanning_mode: u8) -> Vec<u8> {
671        let mut section = vec![0u8; 72];
672        section[..4].copy_from_slice(&(72u32).to_be_bytes());
673        section[4] = 3;
674        section[6..10].copy_from_slice(&(ni * nj).to_be_bytes());
675        section[12..14].copy_from_slice(&0u16.to_be_bytes());
676        section[30..34].copy_from_slice(&ni.to_be_bytes());
677        section[34..38].copy_from_slice(&nj.to_be_bytes());
678        section[46..50].copy_from_slice(&grib_i32_bytes(50_000_000));
679        section[50..54].copy_from_slice(&grib_i32_bytes(-120_000_000));
680        section[55..59].copy_from_slice(&grib_i32_bytes(49_000_000));
681        section[59..63].copy_from_slice(&grib_i32_bytes(-119_000_000));
682        section[63..67].copy_from_slice(&1_000_000u32.to_be_bytes());
683        section[67..71].copy_from_slice(&1_000_000u32.to_be_bytes());
684        section[71] = scanning_mode;
685        section
686    }
687
688    fn build_product(parameter_category: u8, parameter_number: u8) -> Vec<u8> {
689        let mut section = vec![0u8; 34];
690        section[..4].copy_from_slice(&(34u32).to_be_bytes());
691        section[4] = 4;
692        section[7..9].copy_from_slice(&0u16.to_be_bytes());
693        section[9] = parameter_category;
694        section[10] = parameter_number;
695        section[11] = 2;
696        section[17] = 1;
697        section[18..22].copy_from_slice(&0u32.to_be_bytes());
698        section[22] = 103;
699        section[23] = 0;
700        section[24..28].copy_from_slice(&850u32.to_be_bytes());
701        section[28] = 255;
702        section
703    }
704
705    fn build_simple_representation(encoded_values: usize, bits_per_value: u8) -> Vec<u8> {
706        let mut section = vec![0u8; 21];
707        section[..4].copy_from_slice(&(21u32).to_be_bytes());
708        section[4] = 5;
709        section[5..9].copy_from_slice(&(encoded_values as u32).to_be_bytes());
710        section[9..11].copy_from_slice(&0u16.to_be_bytes());
711        section[11..15].copy_from_slice(&0f32.to_be_bytes());
712        section[19] = bits_per_value;
713        section[20] = 0;
714        section
715    }
716
717    fn pack_u8_values(values: &[u8]) -> Vec<u8> {
718        values.to_vec()
719    }
720
721    fn build_bitmap(bits: &[bool]) -> Vec<u8> {
722        let payload_len = bits.len().div_ceil(8) + 1;
723        let mut section = vec![0u8; payload_len + 5];
724        section[..4].copy_from_slice(&((payload_len + 5) as u32).to_be_bytes());
725        section[4] = 6;
726        section[5] = 0;
727        for (index, bit) in bits.iter().copied().enumerate() {
728            if bit {
729                section[6 + index / 8] |= 1 << (7 - (index % 8));
730            }
731        }
732        section
733    }
734
735    fn build_data(payload: &[u8]) -> Vec<u8> {
736        let mut section = vec![0u8; payload.len() + 5];
737        section[..4].copy_from_slice(&((payload.len() + 5) as u32).to_be_bytes());
738        section[4] = 7;
739        section[5..].copy_from_slice(payload);
740        section
741    }
742
743    fn assemble_grib2_message(sections: &[Vec<u8>]) -> Vec<u8> {
744        let total_len = 16 + sections.iter().map(|section| section.len()).sum::<usize>() + 4;
745        let mut message = build_indicator(total_len, 0);
746        for section in sections {
747            message.extend_from_slice(section);
748        }
749        message.extend_from_slice(b"7777");
750        message
751    }
752
753    fn build_grib1_message(values: &[u8]) -> Vec<u8> {
754        let mut pds = vec![0u8; 28];
755        pds[..3].copy_from_slice(&[0, 0, 28]);
756        pds[3] = 2;
757        pds[4] = 7;
758        pds[5] = 255;
759        pds[6] = 0;
760        pds[7] = 0b1000_0000;
761        pds[8] = 11;
762        pds[9] = 100;
763        pds[10..12].copy_from_slice(&850u16.to_be_bytes());
764        pds[12] = 26;
765        pds[13] = 3;
766        pds[14] = 20;
767        pds[15] = 12;
768        pds[16] = 0;
769        pds[17] = 1;
770        pds[18] = 0;
771        pds[19] = 0;
772        pds[20] = 0;
773        pds[24] = 21;
774        pds[25] = 0;
775
776        let mut gds = vec![0u8; 32];
777        gds[..3].copy_from_slice(&[0, 0, 32]);
778        gds[5] = 0;
779        gds[6..8].copy_from_slice(&2u16.to_be_bytes());
780        gds[8..10].copy_from_slice(&2u16.to_be_bytes());
781        gds[10..13].copy_from_slice(&[0x01, 0x4d, 0x50]);
782        gds[13..16].copy_from_slice(&[0x81, 0xd4, 0xc0]);
783        gds[16] = 0x80;
784        gds[17..20].copy_from_slice(&[0x01, 0x49, 0x68]);
785        gds[20..23].copy_from_slice(&[0x81, 0xd0, 0xd8]);
786        gds[23..25].copy_from_slice(&1000u16.to_be_bytes());
787        gds[25..27].copy_from_slice(&1000u16.to_be_bytes());
788
789        let mut bds = vec![0u8; 11 + values.len()];
790        let len = bds.len() as u32;
791        bds[..3].copy_from_slice(&[(len >> 16) as u8, (len >> 8) as u8, len as u8]);
792        bds[3] = 0;
793        bds[10] = 8;
794        bds[11..].copy_from_slice(values);
795
796        let total_len = 8 + pds.len() + gds.len() + bds.len() + 4;
797        let mut message = Vec::new();
798        message.extend_from_slice(b"GRIB");
799        message.extend_from_slice(&[
800            ((total_len >> 16) & 0xff) as u8,
801            ((total_len >> 8) & 0xff) as u8,
802            (total_len & 0xff) as u8,
803            1,
804        ]);
805        message.extend_from_slice(&pds);
806        message.extend_from_slice(&gds);
807        message.extend_from_slice(&bds);
808        message.extend_from_slice(b"7777");
809        message
810    }
811
812    #[test]
813    fn scans_single_grib2_message() {
814        let message = assemble_grib2_message(&[
815            build_identification(),
816            build_grid(2, 2, 0),
817            build_product(0, 0),
818            build_simple_representation(4, 8),
819            build_data(&pack_u8_values(&[1, 2, 3, 4])),
820        ]);
821        let messages = scan_messages(&message, OpenOptions::default()).unwrap();
822        assert_eq!(messages.len(), 1);
823        assert_eq!(messages[0].metadata.parameter.short_name, "TMP");
824    }
825
826    #[test]
827    fn decodes_simple_grib2_message_to_ndarray() {
828        let message = assemble_grib2_message(&[
829            build_identification(),
830            build_grid(2, 2, 0),
831            build_product(0, 0),
832            build_simple_representation(4, 8),
833            build_data(&pack_u8_values(&[1, 2, 3, 4])),
834        ]);
835        let file = GribFile::from_bytes(message).unwrap();
836        let array = file.message(0).unwrap().read_data_as_f64().unwrap();
837        assert_eq!(array.shape(), &[2, 2]);
838        assert_eq!(
839            array.iter().copied().collect::<Vec<_>>(),
840            vec![1.0, 2.0, 3.0, 4.0]
841        );
842    }
843
844    #[test]
845    fn applies_bitmap_to_missing_values() {
846        let message = assemble_grib2_message(&[
847            build_identification(),
848            build_grid(2, 2, 0),
849            build_product(0, 1),
850            build_simple_representation(3, 8),
851            build_bitmap(&[true, false, true, true]),
852            build_data(&pack_u8_values(&[10, 20, 30])),
853        ]);
854        let file = GribFile::from_bytes(message).unwrap();
855        let array = file.message(0).unwrap().read_data_as_f64().unwrap();
856        let values = array.iter().copied().collect::<Vec<_>>();
857        assert_eq!(values[0], 10.0);
858        assert!(values[1].is_nan());
859        assert_eq!(values[2], 20.0);
860        assert_eq!(values[3], 30.0);
861    }
862
863    #[test]
864    fn indexes_multiple_fields_in_one_grib2_message() {
865        let message = assemble_grib2_message(&[
866            build_identification(),
867            build_grid(2, 2, 0),
868            build_product(0, 0),
869            build_simple_representation(4, 8),
870            build_data(&pack_u8_values(&[1, 2, 3, 4])),
871            build_product(0, 2),
872            build_simple_representation(4, 8),
873            build_data(&pack_u8_values(&[5, 6, 7, 8])),
874        ]);
875        let file = GribFile::from_bytes(message).unwrap();
876        assert_eq!(file.message_count(), 2);
877        assert_eq!(file.message(0).unwrap().parameter_name(), "TMP");
878        assert_eq!(file.message(1).unwrap().parameter_name(), "POT");
879    }
880
881    #[test]
882    fn decodes_simple_grib1_message_to_ndarray() {
883        let message = build_grib1_message(&[1, 2, 3, 4]);
884        let file = GribFile::from_bytes(message).unwrap();
885        assert_eq!(file.edition(), 1);
886        let field = file.message(0).unwrap();
887        assert_eq!(field.parameter_name(), "TMP");
888        assert!(field.identification().is_none());
889        assert!(field.grib1_product_definition().is_some());
890        let array = field.read_data_as_f64().unwrap();
891        assert_eq!(array.shape(), &[2, 2]);
892        assert_eq!(
893            array.iter().copied().collect::<Vec<_>>(),
894            vec![1.0, 2.0, 3.0, 4.0]
895        );
896    }
897}