1use bit_vec::BitVec;
2use duckdb::types::{FromSql, FromSqlError, FromSqlResult, ToSql, ToSqlOutput, Value, ValueRef};
3use std::borrow::Cow;
4use std::fmt;
5
6#[derive(Debug, Clone)]
7pub struct Bitstring<'a>(Cow<'a, BitVec>);
12
13impl ToSql for Bitstring<'_> {
14 fn to_sql(&self) -> duckdb::Result<ToSqlOutput<'_>> {
15 if self.as_bitvec().is_empty() {
16 Err(duckdb::Error::ToSqlConversionFailure(Box::new(
17 BitstringError::EmptyBitstring,
18 )))
19 } else {
20 Ok(ToSqlOutput::Owned(Value::Text(format!("{self}"))))
21 }
22 }
23}
24
25impl fmt::Display for Bitstring<'_> {
26 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
27 write!(f, "{}", self.as_bitvec())
28 }
29}
30
31#[derive(Debug, Clone)]
32pub enum BitstringError {
33 EmptyBitstring,
35 RawDataBadPadding(u8),
38 RawDataTooShort(usize),
41}
42
43impl fmt::Display for BitstringError {
44 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
45 match self {
46 Self::RawDataBadPadding(pad) => write!(f, "raw data padding byte should be 0-7, was {pad}"),
47 Self::RawDataTooShort(len) => write!(f, "raw data too short (should be at least 2 bytes, was {len} bytes long)"),
48 Self::EmptyBitstring => write!(f, "DuckDB does not support empty bit strings, consider using a nullable column and Option<Bitstring>")
49 }
50 }
51}
52
53impl std::error::Error for BitstringError {}
54
55impl<'a> Bitstring<'a> {
56 #[must_use]
57 pub fn into_bitvec(self) -> BitVec {
58 self.0.into_owned()
59 }
60
61 #[must_use]
62 pub fn as_bitvec(&'a self) -> &'a BitVec {
63 self.0.as_ref()
64 }
65
66 fn from_raw<'b>(bytes: &[u8]) -> Result<Bitstring<'b>, BitstringError> {
67 if bytes.len() < 2 {
68 Err(BitstringError::RawDataTooShort(bytes.len()))
69 } else if bytes[0] > 7 {
70 Err(BitstringError::RawDataBadPadding(bytes[0]))
71 } else {
72 let mut raw_vec = BitVec::from_bytes(&bytes[1..]);
73 if bytes[0] == 0 {
74 Ok(Bitstring::from(raw_vec))
75 } else {
76 Ok(Bitstring::from(raw_vec.split_off(bytes[0].into())))
77 }
78 }
79 }
80}
81
82impl From<BitVec> for Bitstring<'_> {
83 fn from(v: BitVec) -> Bitstring<'static> {
84 Bitstring(Cow::Owned(v))
85 }
86}
87
88impl<'a> From<&'a BitVec> for Bitstring<'a> {
89 fn from(v: &'a BitVec) -> Self {
90 Self(Cow::Borrowed(v))
91 }
92}
93
94impl FromSql for Bitstring<'_> {
95 fn column_result(value: ValueRef<'_>) -> FromSqlResult<Self> {
96 match value {
97 ValueRef::Blob(bytes) => Ok(Bitstring::from_raw(bytes)?),
98 _ => Err(FromSqlError::InvalidType),
99 }
100 }
101}
102
103impl From<BitstringError> for FromSqlError {
104 fn from(value: BitstringError) -> Self {
105 Self::Other(Box::new(value))
106 }
107}
108
109#[doc = include_str!("../README.md")]
110#[cfg(doctest)]
111pub struct _ReadmeDoctests;
112
113#[cfg(test)]
114mod tests {
115 use crate::{Bitstring, BitstringError};
116 use bit_vec::BitVec;
117 use duckdb::{
118 types::{ToSqlOutput, Value},
119 ToSql,
120 };
121
122 #[test]
123 fn from_raw_1() {
124 let bytes = vec![0, 0b01100101, 0b11100101, 0b00000101];
125 let bv = Bitstring::from_raw(&bytes).unwrap().into_bitvec();
126 let s = format!("{}", bv);
127 assert_eq!(s, "011001011110010100000101");
128 }
129
130 #[test]
131 fn from_raw_2() {
132 let bytes = vec![1, 0b11100101, 0b11100101, 0b00000101];
133 let bv = Bitstring::from_raw(&bytes).unwrap().into_bitvec();
134 let s = format!("{}", bv);
135 assert_eq!(s, "11001011110010100000101");
136 }
137
138 #[test]
139 fn from_raw_3() {
140 let bytes = vec![2, 0b11100101, 0b11100101, 0b00000101];
141 let bv = Bitstring::from_raw(&bytes).unwrap().into_bitvec();
142 let s = format!("{}", bv);
143 assert_eq!(s, "1001011110010100000101");
144 }
145
146 #[test]
147 fn from_raw_4() {
148 let bytes = vec![3, 0b11100101, 0b11100101, 0b00000101];
149 let bv = Bitstring::from_raw(&bytes).unwrap().into_bitvec();
150 let s = format!("{}", bv);
151 assert_eq!(s, "001011110010100000101");
152 }
153
154 #[test]
155 fn from_raw_5() {
156 let bytes = vec![4, 0b11110101, 0b11100101, 0b00000101];
157 let bv = Bitstring::from_raw(&bytes).unwrap().into_bitvec();
158 let s = format!("{}", bv);
159 assert_eq!(s, "01011110010100000101");
160 }
161
162 #[test]
163 fn from_raw_6() {
164 let bytes = vec![5, 0b11111101, 0b11100101, 0b00000101];
165 let bv = Bitstring::from_raw(&bytes).unwrap().into_bitvec();
166 let s = format!("{}", bv);
167 assert_eq!(s, "1011110010100000101");
168 }
169
170 #[test]
171 fn from_raw_7() {
172 let bytes = vec![6, 0b11111101, 0b11100101, 0b00000101];
173 let bv = Bitstring::from_raw(&bytes).unwrap().into_bitvec();
174 let s = format!("{}", bv);
175 assert_eq!(s, "011110010100000101");
176 }
177
178 #[test]
179 fn from_raw_8() {
180 let bytes = vec![7, 0b11111111, 0b11100101, 0b00000101];
181 let bv = Bitstring::from_raw(&bytes).unwrap().into_bitvec();
182 let s = format!("{}", bv);
183 assert_eq!(s, "11110010100000101");
184 }
185
186 #[test]
187 fn from_raw_error_1() {
188 let bytes = vec![8, 0b11111111, 0b11100101, 0b00000101];
189 let bv = Bitstring::from_raw(&bytes);
190 assert!(matches!(bv, Err(BitstringError::RawDataBadPadding(8))));
191 }
192
193 #[test]
194 fn from_raw_error_2() {
195 let bytes = vec![7];
196 let bv = Bitstring::from_raw(&bytes);
197 assert!(matches!(bv, Err(BitstringError::RawDataTooShort(1))));
198 }
199
200 #[test]
201 fn from_raw_error_3() {
202 let bytes = vec![];
203 let bv = Bitstring::from_raw(&bytes);
204 assert!(matches!(bv, Err(BitstringError::RawDataTooShort(0))));
205 }
206
207 #[test]
208 fn test_raw_minimal() {
209 let bytes = vec![7, 0b11111111];
210 let bv = Bitstring::from_raw(&bytes).unwrap().into_bitvec();
211 let s = format!("{}", bv);
212 assert_eq!(s, "1");
213 }
214
215 #[test]
216 fn test_tosql() {
217 let bv = Bitstring::from(BitVec::from_bytes(&[0b11100101, 0b11100101, 0b00000101]));
218 let s = bv.to_sql().unwrap();
219 assert_eq!(
220 s,
221 ToSqlOutput::Owned(Value::Text(String::from("111001011110010100000101")))
222 );
223 }
224
225 #[test]
226 fn test_display() {
227 let bv = Bitstring::from(BitVec::from_bytes(&[0b11100101, 0b11100101, 0b00000101]));
228 let s = format!("{}", bv);
229 assert_eq!(s, "111001011110010100000101");
230 }
231}