1use std::fmt;
12use std::io;
13use std::mem;
14
15use bitvec::prelude::*;
16use libc::pid_t;
17use nix::unistd::Pid;
18use serde::Deserialize;
19use serde::Serialize;
20
21fn longest_run(sequence: &BitVec) -> (usize, usize) {
23 let mut prev_bit = false;
24 let mut prev_count = 1;
25
26 let mut max_count = 0;
27 let mut max_start = 0;
28
29 for (index, bit) in sequence.iter().enumerate() {
30 let count = if index > 0 && prev_bit == *bit {
31 prev_count + 1
32 } else {
33 1
34 };
35 if count > max_count {
36 max_count = count;
37 max_start = index + 1 - count;
38 }
39
40 prev_count = count;
41 prev_bit = *bit;
42 }
43
44 (max_start, max_count)
45}
46
47#[derive(Debug, Clone, Serialize, Deserialize, Default)]
48pub struct Pedigree {
60 pedigree: BitVec,
62}
63
64impl Pedigree {
65 pub fn new() -> Self {
67 Pedigree {
68 pedigree: bitvec![0], }
70 }
71
72 pub fn fork(&self) -> (Self, Self) {
76 let mut parent = self.clone();
77 let child = parent.fork_mut();
78 (parent, child)
79 }
80
81 pub fn fork_mut(&mut self) -> Self {
87 let mut child_pedigree = self.pedigree.clone();
88 child_pedigree.push(true);
89 self.pedigree.push(false);
90 Pedigree {
91 pedigree: child_pedigree,
92 }
93 }
94
95 pub fn raw(&self) -> BitVec {
97 self.pedigree.clone()
98 }
99}
100
101impl TryFrom<&Pedigree> for i32 {
103 type Error = io::Error;
104 fn try_from(pedigree: &Pedigree) -> Result<Self, Self::Error> {
105 const MSB_ZERO_BITS: usize = 1;
107 const TREE_BITS: usize = 16;
108 const RUN_INDEX_BITS: usize = 4;
109 const RUN_TYPE_BITS: usize = 1;
110 const RUN_LENGTH_BITS: usize = 10;
111 debug_assert!(
112 MSB_ZERO_BITS + TREE_BITS + RUN_INDEX_BITS + RUN_TYPE_BITS + RUN_LENGTH_BITS
113 == mem::size_of::<pid_t>() * 8
114 );
115
116 let mut sequence = pedigree.raw();
119 while sequence.len() > 1 && sequence.last().as_deref() == Some(&false) {
120 sequence.pop();
121 }
122
123 let (index, len) = longest_run(&sequence);
125
126 if index >= 2_usize.pow(RUN_INDEX_BITS as u32)
128 || len >= 2_usize.pow(RUN_LENGTH_BITS as u32)
129 || sequence.len() - len > TREE_BITS
130 {
131 Err(Self::Error::other(
132 "Pedigree is too large or complex to be deterministically converted into virtual PID.",
133 ))
134 } else {
135 let mut lower_tree = sequence.split_off(index + len);
137 let run = sequence.split_off(index);
138 let mut tree = sequence;
139 tree.append(&mut lower_tree);
140
141 let mut vpid_bits: BitVec<u32, Msb0> =
143 BitVec::with_capacity(mem::size_of::<pid_t>() * 8);
144
145 vpid_bits.push(false);
148
149 let mut tree_bits: BitVec<u32, Msb0> = BitVec::repeat(false, TREE_BITS - tree.len());
153 tree_bits.append(&mut tree);
154 debug_assert!(tree_bits.len() == TREE_BITS);
155 vpid_bits.append(&mut tree_bits);
156
157 let mut run_index_bits = BitVec::<u32, Msb0>::from_element(index as u32);
158 run_index_bits = run_index_bits.split_off(run_index_bits.len() - RUN_INDEX_BITS);
159 debug_assert!(run_index_bits.len() == RUN_INDEX_BITS);
160 vpid_bits.append(&mut run_index_bits);
161
162 let mut run_type_bits: BitVec<u32, Msb0> = BitVec::new();
163 run_type_bits.push(run[0]);
164 debug_assert!(run_type_bits.len() == RUN_TYPE_BITS);
165 vpid_bits.append(&mut run_type_bits);
166
167 let mut run_length_bits = BitVec::<u32, Msb0>::from_element(len as u32);
168 run_length_bits = run_length_bits.split_off(run_length_bits.len() - RUN_LENGTH_BITS);
169 debug_assert!(run_length_bits.len() == RUN_LENGTH_BITS);
170 vpid_bits.append(&mut run_length_bits);
171
172 debug_assert!(vpid_bits.len() == mem::size_of::<pid_t>() * 8);
173
174 Ok(vpid_bits.into_vec()[0] as i32)
175 }
176 }
177}
178
179impl TryFrom<&Pedigree> for Pid {
180 type Error = io::Error;
181 fn try_from(pedigree: &Pedigree) -> Result<Self, Self::Error> {
182 let num: i32 = i32::try_from(pedigree)?;
183 Ok(Pid::from_raw(num))
184 }
185}
186
187impl fmt::Display for Pedigree {
188 fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), std::fmt::Error> {
189 let len = self.pedigree.len();
190 debug_assert!(len > 0);
192 debug_assert!(!self.pedigree[0]);
193 let mut iter = self.pedigree.iter();
194 let _ = iter.next();
195 for bit in iter {
196 write!(f, "{}", if *bit { "C" } else { "P" })?;
197 }
198 Ok(())
199 }
200}
201
202#[cfg(test)]
203mod test {
204 use super::*;
205
206 #[test]
207 fn test_longest_run() {
208 let sequence = bitvec![1, 1, 0, 0, 0];
209 let (index, len) = longest_run(&sequence);
210 assert_eq!((index, len), (2, 3));
211
212 let sequence = bitvec![1, 1, 1, 0, 0];
213 let (index, len) = longest_run(&sequence);
214 assert_eq!((index, len), (0, 3));
215
216 let sequence = bitvec![1, 1, 0, 0, 1, 1];
217 let (index, len) = longest_run(&sequence);
218 assert_eq!((index, len), (0, 2));
219
220 let sequence = bitvec![1, 0, 0, 0, 0, 1];
221 let (index, len) = longest_run(&sequence);
222 assert_eq!((index, len), (1, 4));
223
224 let sequence = bitvec![1, 0, 1, 0, 1, 0];
225 let (index, len) = longest_run(&sequence);
226 assert_eq!((index, len), (0, 1));
227 }
228
229 #[test]
230 fn test_pedigree_basic() {
231 let mut parent = Pedigree::new();
236 assert_eq!(Pid::try_from(&parent).unwrap(), Pid::from_raw(0x1));
238
239 let child = parent.fork_mut();
240 assert_eq!(Pid::try_from(&parent).unwrap(), Pid::from_raw(0x1));
242 assert_eq!(Pid::try_from(&child).unwrap(), Pid::from_raw(0x00008001));
244
245 let child2 = parent.fork_mut();
246 assert_eq!(Pid::try_from(&parent).unwrap(), Pid::from_raw(0x1));
248 assert_eq!(Pid::try_from(&child2).unwrap(), Pid::from_raw(0x00008002));
250 }
251
252 #[test]
253 fn test_pedigree_many_forks() {
254 let mut many_forks_bitstring = BitVec::repeat(false, 1023);
255 many_forks_bitstring.push(true);
256 let many_forks_pedigree = Pedigree {
257 pedigree: many_forks_bitstring,
258 };
259 assert_eq!(
260 Pid::try_from(&many_forks_pedigree).unwrap(),
261 Pid::from_raw(0x000083FF)
262 );
263 }
264
265 #[test]
266 fn test_pedigree_overflow() {
267 let mut many_forks_bitstring = BitVec::repeat(false, 1024);
268 many_forks_bitstring.push(true);
269 let many_forks_pedigree = Pedigree {
270 pedigree: many_forks_bitstring,
271 };
272 assert!(Pid::try_from(&many_forks_pedigree).is_err());
273 }
274
275 #[test]
276 fn display_pedigree() {
277 let p0 = Pedigree::new();
278 let (p1, _c1) = p0.fork();
279 let (p2, _c2) = p1.fork();
280 let (p3, c3) = p2.fork();
281 assert_eq!(format!("{}", p3), "PPP");
282 assert_eq!(format!("{}", c3), "PPC");
283 }
284}