1use std::os::raw::{c_int, c_void};
3
4use fitsio::FitsFile;
5use fitsio::hdu::HduInfo;
6use fitsio::images::ImageType;
7use fitsio::sys::{self, LONGLONG};
8
9use crate::error::{AtfitsError, Result, ffcall};
10use crate::raw::read_flat;
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum PixelType {
15 F32,
16 F64,
17}
18
19impl PixelType {
20 pub fn from_bitpix(bitpix: i64) -> Self {
22 if bitpix == -64 {
23 PixelType::F64
24 } else {
25 PixelType::F32
26 }
27 }
28}
29
30pub trait CubeElem: Copy + Send + Sync + 'static {
36 const BITPIX: i64;
38 fn read_full(fptr: &mut FitsFile) -> Result<Vec<Self>>;
40 fn read_section(fptr: &mut FitsFile, start: usize, end: usize) -> Result<Vec<Self>>;
42 fn read_section_into(fptr: &mut FitsFile, start: usize, out: &mut [Self]) -> Result<()> {
44 let data = Self::read_section(fptr, start, start + out.len())?;
45 out.copy_from_slice(&data);
46 Ok(())
47 }
48 fn write_section(fptr: &mut FitsFile, start: usize, end: usize, data: &[Self]) -> Result<()>;
50}
51
52pub fn read_plane<T: CubeElem>(
54 fptr: &mut FitsFile,
55 plane_len: usize,
56 index: usize,
57) -> Result<Vec<T>> {
58 let (start, end) = plane_range(plane_len, index)?;
59 T::read_section(fptr, start, end)
60}
61
62pub fn read_plane_into<T: CubeElem>(
65 fptr: &mut FitsFile,
66 index: usize,
67 out: &mut [T],
68) -> Result<()> {
69 let (start, _) = plane_range(out.len(), index)?;
70 T::read_section_into(fptr, start, out)
71}
72
73pub fn write_plane<T: CubeElem>(fptr: &mut FitsFile, index: usize, data: &[T]) -> Result<()> {
75 let (start, end) = plane_range(data.len(), index)?;
76 T::write_section(fptr, start, end, data)
77}
78
79fn plane_range(plane_len: usize, index: usize) -> Result<(usize, usize)> {
80 plane_len
81 .checked_mul(index)
82 .and_then(|start| Some((start, start.checked_add(plane_len)?)))
83 .ok_or_else(|| AtfitsError::Other(format!("plane {index} overflows usize")))
84}
85
86fn primary_image(fptr: &mut FitsFile) -> Result<(usize, bool)> {
88 match fptr.primary_hdu()?.info {
89 HduInfo::ImageInfo { shape, image_type } => {
90 let npix = if shape.is_empty() {
91 0
92 } else {
93 shape.iter().product()
94 };
95 let integer = !matches!(image_type, ImageType::Float | ImageType::Double);
96 Ok((npix, integer))
97 }
98 _ => Err(AtfitsError::Other("primary HDU is not an image".into())),
99 }
100}
101
102fn check_range(start: usize, end: usize, npix: usize) -> Result<()> {
103 if start > end || end > npix {
104 return Err(range_error(start, end, Some(npix)));
105 }
106 Ok(())
107}
108
109fn range_error(start: usize, end: usize, npix: Option<usize>) -> AtfitsError {
110 match npix {
111 Some(n) => AtfitsError::Other(format!(
112 "pixel range {start}..{end} is outside the image ({n} pixels)"
113 )),
114 None => AtfitsError::Other(format!("pixel range {start}..{end} is reversed")),
115 }
116}
117
118macro_rules! impl_cube_elem {
119 ($t:ty, $bitpix:expr, $code:expr) => {
120 impl CubeElem for $t {
121 const BITPIX: i64 = $bitpix;
122
123 fn read_full(fptr: &mut FitsFile) -> Result<Vec<Self>> {
124 let (npix, _) = primary_image(fptr)?;
125 Self::read_section(fptr, 0, npix)
126 }
127
128 fn read_section(fptr: &mut FitsFile, start: usize, end: usize) -> Result<Vec<Self>> {
129 if start > end {
130 return Err(range_error(start, end, None));
131 }
132 let mut out: Vec<$t> = vec![0.0; end - start];
133 Self::read_section_into(fptr, start, &mut out)?;
134 Ok(out)
135 }
136
137 fn read_section_into(
138 fptr: &mut FitsFile,
139 start: usize,
140 out: &mut [Self],
141 ) -> Result<()> {
142 let (npix, integer) = primary_image(fptr)?;
143 let end = start
144 .checked_add(out.len())
145 .ok_or_else(|| range_error(start, usize::MAX, Some(npix)))?;
146 check_range(start, end, npix)?;
147 read_flat(fptr, start, integer.then_some(<$t>::NAN), out)
149 }
150
151 fn write_section(
152 fptr: &mut FitsFile,
153 start: usize,
154 end: usize,
155 data: &[Self],
156 ) -> Result<()> {
157 let (npix, integer) = primary_image(fptr)?;
158 check_range(start, end, npix)?;
159 if data.len() != end - start {
160 return Err(AtfitsError::Other(format!(
161 "pixel range {start}..{end} needs {} values, got {}",
162 end - start,
163 data.len()
164 )));
165 }
166 if !(integer && data.iter().any(|v| v.is_nan())) {
167 return ffcall(|s| unsafe {
169 sys::ffppr(
170 fptr.as_raw(),
171 $code as c_int,
172 start as LONGLONG + 1,
173 data.len() as LONGLONG,
174 data.as_ptr() as *mut c_void,
175 s,
176 )
177 });
178 }
179 let mut nulval = <$t>::NEG_INFINITY;
181 let mut buf: Vec<$t> = data
182 .iter()
183 .map(|&v| if v.is_nan() { nulval } else { v })
184 .collect();
185 ffcall(|s| unsafe {
186 sys::ffppn(
187 fptr.as_raw(),
188 $code as c_int,
189 start as LONGLONG + 1,
190 buf.len() as LONGLONG,
191 buf.as_mut_ptr() as *mut c_void,
192 &mut nulval as *mut $t as *mut c_void,
193 s,
194 )
195 })
196 }
197 }
198 };
199}
200impl_cube_elem!(f32, -32, sys::TFLOAT);
201impl_cube_elem!(f64, -64, sys::TDOUBLE);