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, GRIB2 Lambert conformal
5//! and polar stereographic metadata and flat decode, simple packing, GRIB2
6//! complex packing with general group splitting, and optional image-backed
7//! GRIB2 packing codecs.
8//!
9//! # Example
10//!
11//! ```no_run
12//! use grib_reader::GribFile;
13//!
14//! let file = GribFile::open("gfs.grib2")?;
15//! println!("messages: {}", file.message_count());
16//!
17//! for msg in file.messages() {
18//!     println!(
19//!         "  {} {:?} {:04}-{:02}-{:02}T{:02}:{:02}:{:02}Z",
20//!         msg.parameter_name(),
21//!         msg.grid_shape(),
22//!         msg.reference_time().year,
23//!         msg.reference_time().month,
24//!         msg.reference_time().day,
25//!         msg.reference_time().hour,
26//!         msg.reference_time().minute,
27//!         msg.reference_time().second,
28//!     );
29//! }
30//!
31//! let data = file.message(0)?.read_data_as_f64()?;
32//! println!("shape: {:?}", data.shape());
33//! # Ok::<(), grib_reader::Error>(())
34//! ```
35
36pub mod data;
37pub mod error;
38pub mod grib1;
39pub mod grid;
40pub mod indicator;
41pub mod metadata;
42pub mod parameter;
43pub mod product;
44pub mod sections;
45
46pub use data::{
47    ComplexPackingParams, DataRepresentation, DecodeSample, ImagePackingParams,
48    Jpeg2000PackingParams, PngPackingParams, SimplePackingParams, SpatialDifferencingParams,
49};
50pub use error::{Error, Result};
51pub use grib1::{BinaryDataSection, GridDescription, ProductDefinition as Grib1ProductDefinition};
52pub use grid::{GridDefinition, LambertConformalGrid, LatLonGrid, PolarStereographicGrid};
53pub use metadata::{ForecastTimeUnit, Parameter, ParameterTableSource, ReferenceTime};
54pub use parameter::{
55    LocalParameterEntry, LocalParameterTable, OwnedLocalParameterEntry, BUILTIN_LOCAL_PARAMETERS,
56    LOCAL_PARAMETER_TABLE_CSV_HEADER,
57};
58pub use product::{
59    AnalysisOrForecastTemplate, EnsembleStatisticalProcessTemplate, FixedSurface, Identification,
60    IndividualEnsembleForecastTemplate, ProductDefinition, ProductDefinitionTemplate,
61    StatisticalProcessTemplate, StatisticalTimeRange,
62};
63
64use std::path::Path;
65use std::sync::Arc;
66
67use memmap2::Mmap;
68use ndarray::{ArrayD, IxDyn};
69
70use crate::data::{
71    bitmap_payload as grib2_bitmap_payload, count_bitmap_present_points, decode_field_into,
72    decode_payload_into,
73};
74use crate::indicator::Indicator;
75use crate::parameter::lookup_parameter_with_local_entries;
76use crate::sections::{index_fields, FieldSections, SectionRef};
77
78#[cfg(feature = "rayon")]
79use rayon::prelude::*;
80
81const GRIB_MAGIC: &[u8; 4] = b"GRIB";
82pub const DEFAULT_MAX_DECODED_POINTS: usize = 25_000_000;
83pub const DEFAULT_MAX_AXIS_POINTS: usize = 1_000_000;
84
85/// Configuration for opening GRIB data.
86#[derive(Debug, Clone, Copy)]
87pub struct OpenOptions {
88    /// When `true`, the first malformed GRIB candidate aborts opening.
89    ///
90    /// When `false`, candidate offsets with invalid framing are skipped and
91    /// scanning continues. Once a candidate has a valid indicator, message
92    /// length, and end marker, any indexing or decoding error is returned.
93    pub strict: bool,
94    /// Maximum grid points a decoded field may contain.
95    ///
96    /// This is checked while indexing and before convenience decode methods
97    /// allocate their output buffer. `None` disables the decoded-point limit.
98    pub max_decoded_points: Option<usize>,
99    /// Maximum points a coordinate-axis helper may allocate.
100    ///
101    /// This applies to latitude, longitude, projected x, and projected y axis
102    /// vectors produced through [`Message`] helpers. `None` disables the axis
103    /// limit.
104    pub max_axis_points: Option<usize>,
105}
106
107impl Default for OpenOptions {
108    fn default() -> Self {
109        Self {
110            strict: true,
111            max_decoded_points: Some(DEFAULT_MAX_DECODED_POINTS),
112            max_axis_points: Some(DEFAULT_MAX_AXIS_POINTS),
113        }
114    }
115}
116
117impl OpenOptions {
118    pub fn with_max_decoded_points(mut self, max: usize) -> Self {
119        self.max_decoded_points = Some(max);
120        self
121    }
122
123    pub fn with_max_axis_points(mut self, max: usize) -> Self {
124        self.max_axis_points = Some(max);
125        self
126    }
127
128    pub fn without_decoded_point_limit(mut self) -> Self {
129        self.max_decoded_points = None;
130        self
131    }
132
133    pub fn without_axis_point_limit(mut self) -> Self {
134        self.max_axis_points = None;
135        self
136    }
137
138    pub fn without_limits(mut self) -> Self {
139        self.max_decoded_points = None;
140        self.max_axis_points = None;
141        self
142    }
143}
144
145/// Caller-supplied GRIB1 predefined bitmap.
146///
147/// GRIB1 predefined bitmap table references are center-defined, not globally
148/// standardized. `bitmap` is the raw bitmap payload: one bit per grid point in
149/// GRIB scan order, with `1` meaning present and `0` meaning missing.
150#[derive(Debug, Clone, Copy)]
151pub struct PredefinedBitmap<'a> {
152    pub center_id: u16,
153    pub subcenter_id: Option<u16>,
154    pub table_reference: u16,
155    pub bitmap: &'a [u8],
156}
157
158/// A parsed GRIB field.
159#[derive(Debug, Clone)]
160pub struct MessageMetadata {
161    pub edition: u8,
162    pub center_id: u16,
163    pub subcenter_id: u16,
164    pub discipline: Option<u8>,
165    pub reference_time: ReferenceTime,
166    pub parameter: Parameter,
167    pub grid: GridDefinition,
168    pub data_representation: DataRepresentation,
169    pub forecast_time_unit: Option<u8>,
170    pub forecast_time: Option<u32>,
171    pub message_offset: u64,
172    pub message_length: u64,
173    pub field_index_in_message: usize,
174    grib1_product: Option<grib1::ProductDefinition>,
175    grib2_identification: Option<Identification>,
176    grib2_product: Option<ProductDefinition>,
177}
178
179/// A GRIB file containing one or more logical fields.
180pub struct GribFile {
181    data: GribData,
182    messages: Vec<MessageIndex>,
183    options: OpenOptions,
184}
185
186#[derive(Clone)]
187struct MessageIndex {
188    offset: usize,
189    length: usize,
190    metadata: MessageMetadata,
191    decode_plan: DecodePlan,
192}
193
194#[derive(Clone)]
195enum DecodePlan {
196    Grib1 {
197        bitmap: Option<Grib1BitmapSource>,
198        data: SectionRef,
199    },
200    Grib2(FieldSections),
201}
202
203#[derive(Clone)]
204enum Grib1BitmapSource {
205    Explicit(SectionRef),
206    Predefined(Arc<[u8]>),
207}
208
209enum GribData {
210    Mmap(Mmap),
211    Bytes(Vec<u8>),
212}
213
214impl GribData {
215    fn as_bytes(&self) -> &[u8] {
216        match self {
217            GribData::Mmap(m) => m,
218            GribData::Bytes(b) => b,
219        }
220    }
221}
222
223impl GribFile {
224    /// Open a GRIB file from disk using memory-mapped I/O.
225    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
226        Self::open_with_options(path, OpenOptions::default())
227    }
228
229    /// Open a GRIB file from disk using explicit decoder options.
230    pub fn open_with_options<P: AsRef<Path>>(path: P, options: OpenOptions) -> Result<Self> {
231        Self::open_with_local_parameters(path, options, &[])
232    }
233
234    /// Open a GRIB file from disk using explicit decoder options and local table entries.
235    pub fn open_with_local_parameters<P: AsRef<Path>>(
236        path: P,
237        options: OpenOptions,
238        local_parameters: &[LocalParameterEntry<'_>],
239    ) -> Result<Self> {
240        Self::open_with_local_parameters_and_grib1_predefined_bitmaps(
241            path,
242            options,
243            local_parameters,
244            &[],
245        )
246    }
247
248    /// Open a GRIB file from disk using caller-supplied GRIB1 predefined bitmaps.
249    pub fn open_with_grib1_predefined_bitmaps<P: AsRef<Path>>(
250        path: P,
251        options: OpenOptions,
252        predefined_bitmaps: &[PredefinedBitmap<'_>],
253    ) -> Result<Self> {
254        Self::open_with_local_parameters_and_grib1_predefined_bitmaps(
255            path,
256            options,
257            &[],
258            predefined_bitmaps,
259        )
260    }
261
262    /// Open a GRIB file from disk using GRIB2 local parameters and GRIB1 predefined bitmaps.
263    pub fn open_with_local_parameters_and_grib1_predefined_bitmaps<P: AsRef<Path>>(
264        path: P,
265        options: OpenOptions,
266        local_parameters: &[LocalParameterEntry<'_>],
267        predefined_bitmaps: &[PredefinedBitmap<'_>],
268    ) -> Result<Self> {
269        let file = std::fs::File::open(path.as_ref())
270            .map_err(|e| Error::Io(e, path.as_ref().display().to_string()))?;
271        let mmap = unsafe { Mmap::map(&file) }
272            .map_err(|e| Error::Io(e, path.as_ref().display().to_string()))?;
273        Self::from_data(
274            GribData::Mmap(mmap),
275            options,
276            local_parameters,
277            predefined_bitmaps,
278        )
279    }
280
281    /// Open a GRIB file from an owned byte buffer.
282    pub fn from_bytes(data: Vec<u8>) -> Result<Self> {
283        Self::from_bytes_with_options(data, OpenOptions::default())
284    }
285
286    /// Open a GRIB file from an owned byte buffer using explicit decoder options.
287    pub fn from_bytes_with_options(data: Vec<u8>, options: OpenOptions) -> Result<Self> {
288        Self::from_bytes_with_local_parameters(data, options, &[])
289    }
290
291    /// Open a GRIB file from bytes using explicit decoder options and local table entries.
292    pub fn from_bytes_with_local_parameters(
293        data: Vec<u8>,
294        options: OpenOptions,
295        local_parameters: &[LocalParameterEntry<'_>],
296    ) -> Result<Self> {
297        Self::from_bytes_with_local_parameters_and_grib1_predefined_bitmaps(
298            data,
299            options,
300            local_parameters,
301            &[],
302        )
303    }
304
305    /// Open a GRIB file from bytes using caller-supplied GRIB1 predefined bitmaps.
306    pub fn from_bytes_with_grib1_predefined_bitmaps(
307        data: Vec<u8>,
308        options: OpenOptions,
309        predefined_bitmaps: &[PredefinedBitmap<'_>],
310    ) -> Result<Self> {
311        Self::from_bytes_with_local_parameters_and_grib1_predefined_bitmaps(
312            data,
313            options,
314            &[],
315            predefined_bitmaps,
316        )
317    }
318
319    /// Open a GRIB file from bytes using GRIB2 local parameters and GRIB1 predefined bitmaps.
320    pub fn from_bytes_with_local_parameters_and_grib1_predefined_bitmaps(
321        data: Vec<u8>,
322        options: OpenOptions,
323        local_parameters: &[LocalParameterEntry<'_>],
324        predefined_bitmaps: &[PredefinedBitmap<'_>],
325    ) -> Result<Self> {
326        Self::from_data(
327            GribData::Bytes(data),
328            options,
329            local_parameters,
330            predefined_bitmaps,
331        )
332    }
333
334    fn from_data(
335        data: GribData,
336        options: OpenOptions,
337        local_parameters: &[LocalParameterEntry<'_>],
338        predefined_bitmaps: &[PredefinedBitmap<'_>],
339    ) -> Result<Self> {
340        let messages = scan_messages(
341            data.as_bytes(),
342            options,
343            local_parameters,
344            predefined_bitmaps,
345        )?;
346        if messages.is_empty() {
347            return Err(Error::NoMessages);
348        }
349        Ok(Self {
350            data,
351            messages,
352            options,
353        })
354    }
355
356    /// Returns the GRIB edition of the first field.
357    pub fn edition(&self) -> u8 {
358        self.messages
359            .first()
360            .map(|message| message.metadata.edition)
361            .unwrap_or(0)
362    }
363
364    /// Returns the number of logical fields in the file.
365    pub fn message_count(&self) -> usize {
366        self.messages.len()
367    }
368
369    /// Access a field by index.
370    pub fn message(&self, index: usize) -> Result<Message<'_>> {
371        let record = self
372            .messages
373            .get(index)
374            .ok_or(Error::MessageNotFound(index))?;
375        let bytes = &self.data.as_bytes()[record.offset..record.offset + record.length];
376        Ok(Message {
377            index,
378            bytes,
379            metadata: &record.metadata,
380            decode_plan: record.decode_plan.clone(),
381            options: self.options,
382        })
383    }
384
385    /// Iterate over all fields.
386    pub fn messages(&self) -> MessageIter<'_> {
387        MessageIter {
388            file: self,
389            index: 0,
390        }
391    }
392
393    /// Decode every field in the file.
394    pub fn read_all_data_as_f64(&self) -> Result<Vec<ArrayD<f64>>> {
395        #[cfg(feature = "rayon")]
396        {
397            (0..self.message_count())
398                .into_par_iter()
399                .map(|index| self.message(index)?.read_data_as_f64())
400                .collect()
401        }
402        #[cfg(not(feature = "rayon"))]
403        {
404            (0..self.message_count())
405                .map(|index| self.message(index)?.read_data_as_f64())
406                .collect()
407        }
408    }
409
410    /// Decode every field in the file as `f32`.
411    pub fn read_all_data_as_f32(&self) -> Result<Vec<ArrayD<f32>>> {
412        #[cfg(feature = "rayon")]
413        {
414            (0..self.message_count())
415                .into_par_iter()
416                .map(|index| self.message(index)?.read_data_as_f32())
417                .collect()
418        }
419        #[cfg(not(feature = "rayon"))]
420        {
421            (0..self.message_count())
422                .map(|index| self.message(index)?.read_data_as_f32())
423                .collect()
424        }
425    }
426}
427
428/// A single logical GRIB field.
429pub struct Message<'a> {
430    bytes: &'a [u8],
431    metadata: &'a MessageMetadata,
432    decode_plan: DecodePlan,
433    options: OpenOptions,
434    index: usize,
435}
436
437impl<'a> Message<'a> {
438    pub fn edition(&self) -> u8 {
439        self.metadata.edition
440    }
441
442    pub fn index(&self) -> usize {
443        self.index
444    }
445
446    pub fn metadata(&self) -> &MessageMetadata {
447        self.metadata
448    }
449
450    pub fn reference_time(&self) -> &ReferenceTime {
451        &self.metadata.reference_time
452    }
453
454    pub fn parameter(&self) -> &Parameter {
455        &self.metadata.parameter
456    }
457
458    pub fn center_id(&self) -> u16 {
459        self.metadata.center_id
460    }
461
462    pub fn subcenter_id(&self) -> u16 {
463        self.metadata.subcenter_id
464    }
465
466    pub fn identification(&self) -> Option<&Identification> {
467        self.metadata.grib2_identification.as_ref()
468    }
469
470    pub fn product_definition(&self) -> Option<&ProductDefinition> {
471        self.metadata.grib2_product.as_ref()
472    }
473
474    pub fn grib1_product_definition(&self) -> Option<&grib1::ProductDefinition> {
475        self.metadata.grib1_product.as_ref()
476    }
477
478    pub fn grid_definition(&self) -> &GridDefinition {
479        &self.metadata.grid
480    }
481
482    pub fn parameter_name(&self) -> &str {
483        self.metadata.parameter.short_name.as_ref()
484    }
485
486    pub fn parameter_description(&self) -> &str {
487        self.metadata.parameter.description.as_ref()
488    }
489
490    pub fn forecast_time_unit(&self) -> Option<u8> {
491        self.metadata.forecast_time_unit
492    }
493
494    pub fn forecast_time_unit_kind(&self) -> Option<ForecastTimeUnit> {
495        let unit = self.metadata.forecast_time_unit?;
496        ForecastTimeUnit::from_edition_and_code(self.metadata.edition, unit)
497    }
498
499    pub fn forecast_time(&self) -> Option<u32> {
500        self.metadata.forecast_time
501    }
502
503    pub fn valid_time(&self) -> Option<ReferenceTime> {
504        if let Some(end) = self
505            .metadata
506            .grib2_product
507            .as_ref()
508            .and_then(ProductDefinition::end_of_overall_time_interval)
509        {
510            return Some(end);
511        }
512
513        let unit = self.metadata.forecast_time_unit?;
514        let lead = self.metadata.forecast_time?;
515        self.metadata
516            .reference_time
517            .checked_add_forecast_time_by_edition(self.metadata.edition, unit, lead)
518    }
519
520    pub fn grid_shape(&self) -> (usize, usize) {
521        self.metadata.grid.shape()
522    }
523
524    pub fn latitudes(&self) -> Result<Option<Vec<f64>>> {
525        self.metadata
526            .grid
527            .as_lat_lon()
528            .map(|grid| grid.latitudes_with_limit(self.options.max_axis_points))
529            .transpose()
530    }
531
532    pub fn longitudes(&self) -> Result<Option<Vec<f64>>> {
533        self.metadata
534            .grid
535            .as_lat_lon()
536            .map(|grid| grid.longitudes_with_limit(self.options.max_axis_points))
537            .transpose()
538    }
539
540    pub fn projected_x_coordinates(&self) -> Result<Option<Vec<f64>>> {
541        self.metadata
542            .grid
543            .projected_x_coordinates_with_limit(self.options.max_axis_points)
544    }
545
546    pub fn projected_y_coordinates(&self) -> Result<Option<Vec<f64>>> {
547        self.metadata
548            .grid
549            .projected_y_coordinates_with_limit(self.options.max_axis_points)
550    }
551
552    pub fn decode_into<T: DecodeSample>(&self, out: &mut [T]) -> Result<()> {
553        self.metadata.grid.validate_supported_scan_order()?;
554        let num_grid_points = self.checked_num_points()?;
555        if out.len() != num_grid_points {
556            return Err(Error::DataLengthMismatch {
557                expected: num_grid_points,
558                actual: out.len(),
559            });
560        }
561
562        match &self.decode_plan {
563            DecodePlan::Grib2(field) => {
564                let data_section = section_bytes(self.bytes, field.data);
565                let bitmap_section = match field.bitmap {
566                    Some(section) => grib2_bitmap_payload(section_bytes(self.bytes, section))?,
567                    None => None,
568                };
569                decode_field_into(
570                    data_section,
571                    &self.metadata.data_representation,
572                    bitmap_section,
573                    num_grid_points,
574                    out,
575                )?
576            }
577            DecodePlan::Grib1 { bitmap, data } => {
578                let bitmap_section = match bitmap {
579                    Some(Grib1BitmapSource::Explicit(section)) => {
580                        grib1::bitmap_payload(section_bytes(self.bytes, *section))?
581                    }
582                    Some(Grib1BitmapSource::Predefined(payload)) => Some(payload.as_ref()),
583                    None => None,
584                };
585                let data_section = section_bytes(self.bytes, *data);
586                if data_section.len() < 11 {
587                    return Err(Error::InvalidSection {
588                        section: 4,
589                        reason: format!("expected at least 11 bytes, got {}", data_section.len()),
590                    });
591                }
592                decode_payload_into(
593                    &data_section[11..],
594                    &self.metadata.data_representation,
595                    bitmap_section,
596                    num_grid_points,
597                    out,
598                )?
599            }
600        }
601
602        self.metadata.grid.reorder_for_ndarray_in_place(out)
603    }
604
605    pub fn read_flat_data_as_f64(&self) -> Result<Vec<f64>> {
606        let num_grid_points = self.checked_num_points()?;
607        let mut decoded = zeroed_vec(num_grid_points, 0.0_f64, "decoded f64 field")?;
608        self.decode_into(&mut decoded)?;
609        Ok(decoded)
610    }
611
612    pub fn read_flat_data_as_f32(&self) -> Result<Vec<f32>> {
613        let num_grid_points = self.checked_num_points()?;
614        let mut decoded = zeroed_vec(num_grid_points, 0.0_f32, "decoded f32 field")?;
615        self.decode_into(&mut decoded)?;
616        Ok(decoded)
617    }
618
619    pub fn read_data_as_f64(&self) -> Result<ArrayD<f64>> {
620        let ordered = self.read_flat_data_as_f64()?;
621        ArrayD::from_shape_vec(IxDyn(&self.metadata.grid.ndarray_shape()), ordered)
622            .map_err(|e| Error::Other(format!("failed to build ndarray from decoded field: {e}")))
623    }
624
625    pub fn read_data_as_f32(&self) -> Result<ArrayD<f32>> {
626        let ordered = self.read_flat_data_as_f32()?;
627        ArrayD::from_shape_vec(IxDyn(&self.metadata.grid.ndarray_shape()), ordered)
628            .map_err(|e| Error::Other(format!("failed to build ndarray from decoded field: {e}")))
629    }
630
631    pub fn raw_bytes(&self) -> &[u8] {
632        self.bytes
633    }
634
635    fn checked_num_points(&self) -> Result<usize> {
636        let num_grid_points = self.metadata.grid.checked_num_points()?;
637        ensure_limit(
638            "decoded grid points",
639            num_grid_points,
640            self.options.max_decoded_points,
641        )?;
642        Ok(num_grid_points)
643    }
644}
645
646/// Iterator over fields in a GRIB file.
647pub struct MessageIter<'a> {
648    file: &'a GribFile,
649    index: usize,
650}
651
652impl<'a> Iterator for MessageIter<'a> {
653    type Item = Message<'a>;
654
655    fn next(&mut self) -> Option<Self::Item> {
656        if self.index >= self.file.message_count() {
657            return None;
658        }
659        let message = self.file.message(self.index).ok()?;
660        self.index += 1;
661        Some(message)
662    }
663}
664
665fn section_bytes(msg_bytes: &[u8], section: SectionRef) -> &[u8] {
666    &msg_bytes[section.offset..section.offset + section.length]
667}
668
669fn zeroed_vec<T: Copy>(len: usize, value: T, what: &'static str) -> Result<Vec<T>> {
670    let mut out = Vec::new();
671    out.try_reserve(len)
672        .map_err(|e| Error::Other(format!("failed to reserve {len} {what} values: {e}")))?;
673    out.resize(len, value);
674    Ok(out)
675}
676
677fn ensure_limit(what: &'static str, requested: usize, limit: Option<usize>) -> Result<()> {
678    if let Some(limit) = limit {
679        if requested > limit {
680            return Err(Error::LimitExceeded {
681                what,
682                requested,
683                limit,
684            });
685        }
686    }
687    Ok(())
688}
689
690fn scan_messages(
691    data: &[u8],
692    options: OpenOptions,
693    local_parameters: &[LocalParameterEntry<'_>],
694    predefined_bitmaps: &[PredefinedBitmap<'_>],
695) -> Result<Vec<MessageIndex>> {
696    let mut messages = Vec::new();
697    let mut pos = 0usize;
698
699    while pos + 8 <= data.len() {
700        if &data[pos..pos + 4] != GRIB_MAGIC {
701            pos += 1;
702            continue;
703        }
704
705        let (indicator, next_pos) = match locate_message(data, pos) {
706            Ok(located) => located,
707            Err(err) if !options.strict && is_recoverable_candidate_error(&err) => {
708                pos += 4;
709                continue;
710            }
711            Err(err) => return Err(err),
712        };
713
714        let message_bytes = &data[pos..next_pos];
715        let indexed = match indicator.edition {
716            1 => index_grib1_message(message_bytes, pos, options, predefined_bitmaps)?,
717            2 => index_grib2_message(message_bytes, pos, options, &indicator, local_parameters)?,
718            other => return Err(Error::UnsupportedEdition(other)),
719        };
720
721        messages.extend(indexed);
722        pos = next_pos;
723    }
724
725    Ok(messages)
726}
727
728fn is_recoverable_candidate_error(err: &Error) -> bool {
729    matches!(err, Error::InvalidMessage(_) | Error::Truncated { .. })
730}
731
732fn locate_message(data: &[u8], pos: usize) -> Result<(Indicator, usize)> {
733    let edition = Indicator::edition(&data[pos..]).ok_or_else(|| {
734        Error::InvalidMessage(format!("failed to parse indicator at byte offset {pos}"))
735    })?;
736    if !matches!(edition, 1 | 2) {
737        return Err(Error::UnsupportedEdition(edition));
738    }
739
740    let indicator = Indicator::parse(&data[pos..]).ok_or_else(|| {
741        Error::InvalidMessage(format!("failed to parse indicator at byte offset {pos}"))
742    })?;
743    let length = indicator.total_length as usize;
744    if length < 12 {
745        return Err(Error::InvalidMessage(format!(
746            "message at byte offset {pos} reports impossible length {length}"
747        )));
748    }
749    let end = pos
750        .checked_add(length)
751        .ok_or_else(|| Error::InvalidMessage("message length overflow".into()))?;
752    if end > data.len() {
753        return Err(Error::Truncated { offset: end as u64 });
754    }
755    if &data[end - 4..end] != b"7777" {
756        return Err(Error::InvalidMessage(format!(
757            "message at byte offset {pos} does not end with 7777"
758        )));
759    }
760
761    Ok((indicator, end))
762}
763
764fn index_grib1_message(
765    message_bytes: &[u8],
766    offset: usize,
767    options: OpenOptions,
768    predefined_bitmaps: &[PredefinedBitmap<'_>],
769) -> Result<Vec<MessageIndex>> {
770    let sections = grib1::parse_message_sections(message_bytes)?;
771    let grid_ref = sections.grid.ok_or_else(|| {
772        Error::InvalidSectionOrder(
773            "GRIB1 decoding requires an explicit grid definition section".into(),
774        )
775    })?;
776    let grid_description = GridDescription::parse(section_bytes(message_bytes, grid_ref))?;
777    let grid = grid_description.grid;
778    let num_grid_points = validate_grid_limits(&grid, options)?;
779
780    let bitmap = match sections.bitmap {
781        Some(section) => {
782            let section_bytes = section_bytes(message_bytes, section);
783            let table_reference = grib1::bitmap_table_reference(section_bytes)?;
784            if table_reference == 0 {
785                Some(Grib1BitmapSource::Explicit(section))
786            } else {
787                let payload = resolve_predefined_bitmap(
788                    predefined_bitmaps,
789                    sections.product.center_id as u16,
790                    sections.product.subcenter_id as u16,
791                    table_reference,
792                )?;
793                Some(Grib1BitmapSource::Predefined(Arc::from(payload)))
794            }
795        }
796        None => None,
797    };
798
799    let bitmap_present_count = match &bitmap {
800        Some(Grib1BitmapSource::Explicit(section)) => {
801            let bitmap_payload = grib1::bitmap_payload(section_bytes(message_bytes, *section))?
802                .ok_or(Error::MissingBitmap)?;
803            count_bitmap_present_points(bitmap_payload, num_grid_points)?
804        }
805        Some(Grib1BitmapSource::Predefined(payload)) => {
806            count_bitmap_present_points(payload, num_grid_points)?
807        }
808        None => num_grid_points,
809    };
810
811    let (_bds, data_representation) = BinaryDataSection::parse(
812        section_bytes(message_bytes, sections.data),
813        sections.product.decimal_scale,
814        bitmap_present_count,
815    )?;
816    let parameter = sections.product.parameter();
817
818    Ok(vec![MessageIndex {
819        offset,
820        length: message_bytes.len(),
821        decode_plan: DecodePlan::Grib1 {
822            bitmap,
823            data: sections.data,
824        },
825        metadata: MessageMetadata {
826            edition: 1,
827            center_id: sections.product.center_id as u16,
828            subcenter_id: sections.product.subcenter_id as u16,
829            discipline: None,
830            reference_time: sections.product.reference_time,
831            parameter,
832            grid,
833            data_representation,
834            forecast_time_unit: Some(sections.product.forecast_time_unit),
835            forecast_time: sections.product.forecast_time(),
836            message_offset: offset as u64,
837            message_length: message_bytes.len() as u64,
838            field_index_in_message: 0,
839            grib1_product: Some(sections.product),
840            grib2_identification: None,
841            grib2_product: None,
842        },
843    }])
844}
845
846fn resolve_predefined_bitmap<'a>(
847    predefined_bitmaps: &[PredefinedBitmap<'a>],
848    center_id: u16,
849    subcenter_id: u16,
850    table_reference: u16,
851) -> Result<&'a [u8]> {
852    predefined_bitmaps
853        .iter()
854        .find(|entry| {
855            entry.center_id == center_id
856                && entry.subcenter_id == Some(subcenter_id)
857                && entry.table_reference == table_reference
858        })
859        .or_else(|| {
860            predefined_bitmaps.iter().find(|entry| {
861                entry.center_id == center_id
862                    && entry.subcenter_id.is_none()
863                    && entry.table_reference == table_reference
864            })
865        })
866        .map(|entry| entry.bitmap)
867        .ok_or(Error::UnsupportedBitmapIndicator(table_reference))
868}
869
870fn validate_grid_limits(grid: &GridDefinition, options: OpenOptions) -> Result<usize> {
871    let num_grid_points = grid.checked_num_points()?;
872    ensure_limit(
873        "decoded grid points",
874        num_grid_points,
875        options.max_decoded_points,
876    )?;
877    Ok(num_grid_points)
878}
879
880fn validate_grib2_grid_points(grid: &GridDefinition) -> Result<()> {
881    let Some(declared) = grid.declared_num_points() else {
882        return Ok(());
883    };
884    let projected = grid.shape_num_points()?;
885    if declared != projected {
886        return Err(Error::InvalidSection {
887            section: 3,
888            reason: format!(
889                "grid template 3.{} declares {declared} points but Nx*Ny is {projected}",
890                grid.template_number()
891            ),
892        });
893    }
894    Ok(())
895}
896
897fn index_grib2_message(
898    message_bytes: &[u8],
899    offset: usize,
900    options: OpenOptions,
901    indicator: &Indicator,
902    local_parameters: &[LocalParameterEntry<'_>],
903) -> Result<Vec<MessageIndex>> {
904    let fields = index_fields(message_bytes)?;
905    let mut messages = Vec::with_capacity(fields.len());
906
907    for (field_index_in_message, field_sections) in fields.into_iter().enumerate() {
908        let identification =
909            Identification::parse(section_bytes(message_bytes, field_sections.identification))?;
910        let grid_section = section_bytes(message_bytes, field_sections.grid);
911        let grid = GridDefinition::parse(grid_section)?;
912        validate_grib2_grid_points(&grid)?;
913        validate_grid_limits(&grid, options)?;
914        let product =
915            ProductDefinition::parse(section_bytes(message_bytes, field_sections.product))?;
916        let data_representation = DataRepresentation::parse(section_bytes(
917            message_bytes,
918            field_sections.data_representation,
919        ))?;
920        let parameter = lookup_parameter_with_local_entries(
921            indicator.discipline,
922            product.parameter_category,
923            product.parameter_number,
924            identification.center_id,
925            identification.subcenter_id,
926            identification.local_table_version,
927            local_parameters,
928        );
929
930        messages.push(MessageIndex {
931            offset,
932            length: message_bytes.len(),
933            decode_plan: DecodePlan::Grib2(field_sections),
934            metadata: MessageMetadata {
935                edition: indicator.edition,
936                center_id: identification.center_id,
937                subcenter_id: identification.subcenter_id,
938                discipline: Some(indicator.discipline),
939                reference_time: ReferenceTime {
940                    year: identification.reference_year,
941                    month: identification.reference_month,
942                    day: identification.reference_day,
943                    hour: identification.reference_hour,
944                    minute: identification.reference_minute,
945                    second: identification.reference_second,
946                },
947                parameter,
948                grid,
949                data_representation,
950                forecast_time_unit: product.forecast_time_unit(),
951                forecast_time: product.forecast_time(),
952                message_offset: offset as u64,
953                message_length: message_bytes.len() as u64,
954                field_index_in_message,
955                grib1_product: None,
956                grib2_identification: Some(identification),
957                grib2_product: Some(product),
958            },
959        });
960    }
961
962    Ok(messages)
963}
964
965#[cfg(test)]
966mod tests {
967    use super::*;
968
969    fn grib_i32_bytes(value: i32) -> [u8; 4] {
970        if value >= 0 {
971            (value as u32).to_be_bytes()
972        } else {
973            ((-value) as u32 | 0x8000_0000).to_be_bytes()
974        }
975    }
976
977    fn build_indicator(total_len: usize, discipline: u8) -> Vec<u8> {
978        let mut indicator = Vec::with_capacity(16);
979        indicator.extend_from_slice(b"GRIB");
980        indicator.extend_from_slice(&[0, 0]);
981        indicator.push(discipline);
982        indicator.push(2);
983        indicator.extend_from_slice(&(total_len as u64).to_be_bytes());
984        indicator
985    }
986
987    fn build_identification() -> Vec<u8> {
988        let mut section = vec![0u8; 21];
989        section[..4].copy_from_slice(&(21u32).to_be_bytes());
990        section[4] = 1;
991        section[5..7].copy_from_slice(&7u16.to_be_bytes());
992        section[7..9].copy_from_slice(&0u16.to_be_bytes());
993        section[9] = 35;
994        section[10] = 1;
995        section[11] = 1;
996        section[12..14].copy_from_slice(&2026u16.to_be_bytes());
997        section[14] = 3;
998        section[15] = 20;
999        section[16] = 12;
1000        section[17] = 0;
1001        section[18] = 0;
1002        section[19] = 0;
1003        section[20] = 1;
1004        section
1005    }
1006
1007    fn build_grid(ni: u32, nj: u32, scanning_mode: u8) -> Vec<u8> {
1008        let mut section = vec![0u8; 72];
1009        section[..4].copy_from_slice(&(72u32).to_be_bytes());
1010        section[4] = 3;
1011        section[6..10].copy_from_slice(&(ni * nj).to_be_bytes());
1012        section[12..14].copy_from_slice(&0u16.to_be_bytes());
1013        section[30..34].copy_from_slice(&ni.to_be_bytes());
1014        section[34..38].copy_from_slice(&nj.to_be_bytes());
1015        section[46..50].copy_from_slice(&grib_i32_bytes(50_000_000));
1016        section[50..54].copy_from_slice(&grib_i32_bytes(-120_000_000));
1017        section[55..59].copy_from_slice(&grib_i32_bytes(49_000_000));
1018        section[59..63].copy_from_slice(&grib_i32_bytes(-119_000_000));
1019        section[63..67].copy_from_slice(&1_000_000u32.to_be_bytes());
1020        section[67..71].copy_from_slice(&1_000_000u32.to_be_bytes());
1021        section[71] = scanning_mode;
1022        section
1023    }
1024
1025    fn build_product(parameter_category: u8, parameter_number: u8) -> Vec<u8> {
1026        let mut section = vec![0u8; 34];
1027        section[..4].copy_from_slice(&(34u32).to_be_bytes());
1028        section[4] = 4;
1029        section[7..9].copy_from_slice(&0u16.to_be_bytes());
1030        section[9] = parameter_category;
1031        section[10] = parameter_number;
1032        section[11] = 2;
1033        section[17] = 1;
1034        section[18..22].copy_from_slice(&0u32.to_be_bytes());
1035        section[22] = 103;
1036        section[23] = 0;
1037        section[24..28].copy_from_slice(&850u32.to_be_bytes());
1038        section[28] = 255;
1039        section
1040    }
1041
1042    fn build_simple_representation(encoded_values: usize, bits_per_value: u8) -> Vec<u8> {
1043        let mut section = vec![0u8; 21];
1044        section[..4].copy_from_slice(&(21u32).to_be_bytes());
1045        section[4] = 5;
1046        section[5..9].copy_from_slice(&(encoded_values as u32).to_be_bytes());
1047        section[9..11].copy_from_slice(&0u16.to_be_bytes());
1048        section[11..15].copy_from_slice(&0f32.to_be_bytes());
1049        section[19] = bits_per_value;
1050        section[20] = 0;
1051        section
1052    }
1053
1054    fn pack_u8_values(values: &[u8]) -> Vec<u8> {
1055        values.to_vec()
1056    }
1057
1058    fn build_bitmap(bits: &[bool]) -> Vec<u8> {
1059        let payload_len = bits.len().div_ceil(8) + 1;
1060        let mut section = vec![0u8; payload_len + 5];
1061        section[..4].copy_from_slice(&((payload_len + 5) as u32).to_be_bytes());
1062        section[4] = 6;
1063        section[5] = 0;
1064        for (index, bit) in bits.iter().copied().enumerate() {
1065            if bit {
1066                section[6 + index / 8] |= 1 << (7 - (index % 8));
1067            }
1068        }
1069        section
1070    }
1071
1072    fn build_data(payload: &[u8]) -> Vec<u8> {
1073        let mut section = vec![0u8; payload.len() + 5];
1074        section[..4].copy_from_slice(&((payload.len() + 5) as u32).to_be_bytes());
1075        section[4] = 7;
1076        section[5..].copy_from_slice(payload);
1077        section
1078    }
1079
1080    fn assemble_grib2_message(sections: &[Vec<u8>]) -> Vec<u8> {
1081        let total_len = 16 + sections.iter().map(|section| section.len()).sum::<usize>() + 4;
1082        let mut message = build_indicator(total_len, 0);
1083        for section in sections {
1084            message.extend_from_slice(section);
1085        }
1086        message.extend_from_slice(b"7777");
1087        message
1088    }
1089
1090    fn build_grib1_message(values: &[u8]) -> Vec<u8> {
1091        let mut pds = vec![0u8; 28];
1092        pds[..3].copy_from_slice(&[0, 0, 28]);
1093        pds[3] = 2;
1094        pds[4] = 7;
1095        pds[5] = 255;
1096        pds[6] = 0;
1097        pds[7] = 0b1000_0000;
1098        pds[8] = 11;
1099        pds[9] = 100;
1100        pds[10..12].copy_from_slice(&850u16.to_be_bytes());
1101        pds[12] = 26;
1102        pds[13] = 3;
1103        pds[14] = 20;
1104        pds[15] = 12;
1105        pds[16] = 0;
1106        pds[17] = 1;
1107        pds[18] = 0;
1108        pds[19] = 0;
1109        pds[20] = 0;
1110        pds[24] = 21;
1111        pds[25] = 0;
1112
1113        let mut gds = vec![0u8; 32];
1114        gds[..3].copy_from_slice(&[0, 0, 32]);
1115        gds[5] = 0;
1116        gds[6..8].copy_from_slice(&2u16.to_be_bytes());
1117        gds[8..10].copy_from_slice(&2u16.to_be_bytes());
1118        gds[10..13].copy_from_slice(&[0x01, 0x4d, 0x50]);
1119        gds[13..16].copy_from_slice(&[0x81, 0xd4, 0xc0]);
1120        gds[16] = 0x80;
1121        gds[17..20].copy_from_slice(&[0x01, 0x49, 0x68]);
1122        gds[20..23].copy_from_slice(&[0x81, 0xd0, 0xd8]);
1123        gds[23..25].copy_from_slice(&1000u16.to_be_bytes());
1124        gds[25..27].copy_from_slice(&1000u16.to_be_bytes());
1125
1126        let mut bds = vec![0u8; 11 + values.len()];
1127        let len = bds.len() as u32;
1128        bds[..3].copy_from_slice(&[(len >> 16) as u8, (len >> 8) as u8, len as u8]);
1129        bds[3] = 0;
1130        bds[10] = 8;
1131        bds[11..].copy_from_slice(values);
1132
1133        let total_len = 8 + pds.len() + gds.len() + bds.len() + 4;
1134        let mut message = Vec::new();
1135        message.extend_from_slice(b"GRIB");
1136        message.extend_from_slice(&[
1137            ((total_len >> 16) & 0xff) as u8,
1138            ((total_len >> 8) & 0xff) as u8,
1139            (total_len & 0xff) as u8,
1140            1,
1141        ]);
1142        message.extend_from_slice(&pds);
1143        message.extend_from_slice(&gds);
1144        message.extend_from_slice(&bds);
1145        message.extend_from_slice(b"7777");
1146        message
1147    }
1148
1149    fn build_grib1_message_with_predefined_bitmap(values: &[u8], table_reference: u16) -> Vec<u8> {
1150        let mut message = build_grib1_message(values);
1151        message[8 + 7] = 0b1100_0000;
1152
1153        let mut bitmap = [0u8; 6];
1154        bitmap[..3].copy_from_slice(&[0, 0, 6]);
1155        bitmap[4..6].copy_from_slice(&table_reference.to_be_bytes());
1156
1157        let bitmap_offset = 8 + 28 + 32;
1158        message.splice(bitmap_offset..bitmap_offset, bitmap);
1159        let total_len = message.len();
1160        message[4] = ((total_len >> 16) & 0xff) as u8;
1161        message[5] = ((total_len >> 8) & 0xff) as u8;
1162        message[6] = (total_len & 0xff) as u8;
1163        message
1164    }
1165
1166    #[test]
1167    fn scans_single_grib2_message() {
1168        let message = assemble_grib2_message(&[
1169            build_identification(),
1170            build_grid(2, 2, 0),
1171            build_product(0, 0),
1172            build_simple_representation(4, 8),
1173            build_data(&pack_u8_values(&[1, 2, 3, 4])),
1174        ]);
1175        let messages = scan_messages(&message, OpenOptions::default(), &[], &[]).unwrap();
1176        assert_eq!(messages.len(), 1);
1177        assert_eq!(messages[0].metadata.parameter.short_name, "TMP");
1178    }
1179
1180    #[test]
1181    fn decodes_simple_grib2_message_to_ndarray() {
1182        let message = assemble_grib2_message(&[
1183            build_identification(),
1184            build_grid(2, 2, 0),
1185            build_product(0, 0),
1186            build_simple_representation(4, 8),
1187            build_data(&pack_u8_values(&[1, 2, 3, 4])),
1188        ]);
1189        let file = GribFile::from_bytes(message).unwrap();
1190        let array = file.message(0).unwrap().read_data_as_f64().unwrap();
1191        assert_eq!(array.shape(), &[2, 2]);
1192        assert_eq!(
1193            array.iter().copied().collect::<Vec<_>>(),
1194            vec![1.0, 2.0, 3.0, 4.0]
1195        );
1196    }
1197
1198    #[test]
1199    fn applies_bitmap_to_missing_values() {
1200        let message = assemble_grib2_message(&[
1201            build_identification(),
1202            build_grid(2, 2, 0),
1203            build_product(0, 1),
1204            build_simple_representation(3, 8),
1205            build_bitmap(&[true, false, true, true]),
1206            build_data(&pack_u8_values(&[10, 20, 30])),
1207        ]);
1208        let file = GribFile::from_bytes(message).unwrap();
1209        let array = file.message(0).unwrap().read_data_as_f64().unwrap();
1210        let values = array.iter().copied().collect::<Vec<_>>();
1211        assert_eq!(values[0], 10.0);
1212        assert!(values[1].is_nan());
1213        assert_eq!(values[2], 20.0);
1214        assert_eq!(values[3], 30.0);
1215    }
1216
1217    #[test]
1218    fn indexes_multiple_fields_in_one_grib2_message() {
1219        let message = assemble_grib2_message(&[
1220            build_identification(),
1221            build_grid(2, 2, 0),
1222            build_product(0, 0),
1223            build_simple_representation(4, 8),
1224            build_data(&pack_u8_values(&[1, 2, 3, 4])),
1225            build_product(0, 2),
1226            build_simple_representation(4, 8),
1227            build_data(&pack_u8_values(&[5, 6, 7, 8])),
1228        ]);
1229        let file = GribFile::from_bytes(message).unwrap();
1230        assert_eq!(file.message_count(), 2);
1231        assert_eq!(file.message(0).unwrap().parameter_name(), "TMP");
1232        assert_eq!(file.message(1).unwrap().parameter_name(), "POT");
1233    }
1234
1235    #[test]
1236    fn decodes_simple_grib1_message_to_ndarray() {
1237        let message = build_grib1_message(&[1, 2, 3, 4]);
1238        let file = GribFile::from_bytes(message).unwrap();
1239        assert_eq!(file.edition(), 1);
1240        let field = file.message(0).unwrap();
1241        assert_eq!(field.parameter_name(), "TMP");
1242        assert!(field.identification().is_none());
1243        assert!(field.grib1_product_definition().is_some());
1244        let array = field.read_data_as_f64().unwrap();
1245        assert_eq!(array.shape(), &[2, 2]);
1246        assert_eq!(
1247            array.iter().copied().collect::<Vec<_>>(),
1248            vec![1.0, 2.0, 3.0, 4.0]
1249        );
1250    }
1251
1252    #[test]
1253    fn decodes_grib1_predefined_bitmap_when_supplied() {
1254        let message = build_grib1_message_with_predefined_bitmap(&[9, 7], 300);
1255        let bitmap = [0b1010_0000];
1256        let predefined = [PredefinedBitmap {
1257            center_id: 7,
1258            subcenter_id: Some(0),
1259            table_reference: 300,
1260            bitmap: &bitmap,
1261        }];
1262
1263        let file = GribFile::from_bytes_with_grib1_predefined_bitmaps(
1264            message,
1265            OpenOptions::default(),
1266            &predefined,
1267        )
1268        .unwrap();
1269        let decoded = file.message(0).unwrap().read_flat_data_as_f64().unwrap();
1270        assert_eq!(decoded[0], 9.0);
1271        assert!(decoded[1].is_nan());
1272        assert_eq!(decoded[2], 7.0);
1273        assert!(decoded[3].is_nan());
1274    }
1275
1276    #[test]
1277    fn rejects_grib1_predefined_bitmap_without_supplied_table() {
1278        let message = build_grib1_message_with_predefined_bitmap(&[9, 7], 300);
1279        let err = match GribFile::from_bytes(message) {
1280            Ok(_) => panic!("expected unsupported predefined bitmap"),
1281            Err(err) => err,
1282        };
1283
1284        assert!(matches!(err, Error::UnsupportedBitmapIndicator(300)));
1285    }
1286}