Fix bug which prevents selection of specific constant chains

Signed-off-by: Douwe Schulte <d.schulte@uu.nl>

rustyms/src/imgt/select.rs

@@ -74,7 +74,11 @@ impl Selection {
             .flat_map(|g| g.into_iter().map(|c| (g.species, c.0, c.1)))
             .filter(move |(_, kind, _)| self.chains.as_ref().map_or(true, |k| k.contains(kind)))
             .flat_map(|(species, _, c)| c.into_iter().map(move |g| (species, g.0, g.1)))
-            .filter(move |(_, gene, _)| self.genes.as_ref().map_or(true, |s| s.contains(gene)))
+            .filter(move |(_, gene, _)| {
+                self.genes
+                    .as_ref()
+                    .map_or(true, |s| contains_gene(s, *gene))
+            })
             .flat_map(|(species, _, germlines)| germlines.iter().map(move |a| (species, a)))
             .flat_map(move |(species, germline)| {
                 germline
@@ -97,7 +101,11 @@ impl Selection {
             .flat_map(|g| g.into_par_iter().map(|c| (g.species, c.0, c.1)))
             .filter(move |(_, kind, _)| self.chains.as_ref().map_or(true, |k| k.contains(kind)))
             .flat_map(|(species, _, c)| c.into_par_iter().map(move |g| (species, g.0, g.1)))
-            .filter(move |(_, gene, _)| self.genes.as_ref().map_or(true, |s| s.contains(gene)))
+            .filter(move |(_, gene, _)| {
+                self.genes
+                    .as_ref()
+                    .map_or(true, |s| contains_gene(s, *gene))
+            })
             .flat_map(|(species, _, germlines)| {
                 germlines.into_par_iter().map(move |a| (species, a))
             })
@@ -111,6 +119,10 @@ impl Selection {
     }
 }
 
+fn contains_gene(s: &HashSet<GeneType>, gene: GeneType) -> bool {
+    s.contains(&gene) || matches!(gene, GeneType::C(_)) && s.contains(&GeneType::C(None))
+}
+
 impl Default for Selection {
     /// Get a default selection, which gives all kinds and genes but only returns the first allele
     fn default() -> Self {
@@ -220,7 +232,14 @@ impl Germlines {
         let genes = match gene.gene {
             GeneType::V => &chain.variable,
             GeneType::J => &chain.joining,
-            GeneType::C(_) => &chain.constant,
+            GeneType::C(None) => &chain.c,
+            GeneType::C(Some(Constant::A)) => &chain.a,
+            GeneType::C(Some(Constant::D)) => &chain.d,
+            GeneType::C(Some(Constant::E)) => &chain.e,
+            GeneType::C(Some(Constant::G)) => &chain.g,
+            GeneType::C(Some(Constant::M)) => &chain.m,
+            GeneType::C(Some(Constant::O)) => &chain.o,
+            GeneType::C(Some(Constant::T)) => &chain.t,
         };
         genes
             .binary_search_by(|g| g.name.cmp(&gene))
@@ -245,6 +264,10 @@ impl Germlines {
 #[cfg(test)]
 #[allow(clippy::missing_panics_doc)]
 mod tests {
+    use std::collections::HashSet;
+
+    use crate::imgt::select::contains_gene;
+
     use super::Selection;
     use super::{ChainType, GeneType, Species};
 
@@ -257,4 +280,38 @@ mod tests {
         let first = selection.germlines().next().unwrap();
         assert_eq!(first.name(), "IGHV1-2*01");
     }
+
+    #[test]
+    fn try_first_g_human() {
+        let selection = Selection::default()
+            .species([Species::HomoSapiens])
+            .chain([ChainType::Heavy])
+            .gene([GeneType::C(Some(crate::imgt::Constant::G))]);
+        let first = selection.germlines().next().unwrap();
+        assert_eq!(first.name(), "IGG1");
+    }
+
+    #[test]
+    fn gene_selections() {
+        let constant = HashSet::from([GeneType::C(None)]);
+        assert!(contains_gene(&constant, GeneType::C(None)));
+        assert!(contains_gene(
+            &constant,
+            GeneType::C(Some(crate::imgt::Constant::G))
+        ));
+        assert!(contains_gene(
+            &constant,
+            GeneType::C(Some(crate::imgt::Constant::A))
+        ));
+        let constant_g = HashSet::from([GeneType::C(Some(crate::imgt::Constant::G))]);
+        assert!(!contains_gene(&constant_g, GeneType::C(None)));
+        assert!(contains_gene(
+            &constant_g,
+            GeneType::C(Some(crate::imgt::Constant::G))
+        ));
+        assert!(!contains_gene(
+            &constant_g,
+            GeneType::C(Some(crate::imgt::Constant::A))
+        ));
+    }
 }

rustyms/src/imgt/shared/regions.rs

@@ -73,7 +73,14 @@ impl<'a> IntoParallelIterator for &'a Germlines {
 pub(crate) struct Chain {
     pub variable: Vec<Germline>,
     pub joining: Vec<Germline>,
-    pub constant: Vec<Germline>,
+    pub c: Vec<Germline>,
+    pub a: Vec<Germline>,
+    pub d: Vec<Germline>,
+    pub e: Vec<Germline>,
+    pub g: Vec<Germline>,
+    pub m: Vec<Germline>,
+    pub o: Vec<Germline>,
+    pub t: Vec<Germline>,
 }
 
 impl Chain {
@@ -83,7 +90,14 @@ impl Chain {
         let db = match &germline.name.gene {
             GeneType::V => &mut self.variable,
             GeneType::J => &mut self.joining,
-            GeneType::C(_) => &mut self.constant,
+            GeneType::C(None) => &mut self.c,
+            GeneType::C(Some(Constant::A)) => &mut self.a,
+            GeneType::C(Some(Constant::D)) => &mut self.d,
+            GeneType::C(Some(Constant::E)) => &mut self.e,
+            GeneType::C(Some(Constant::G)) => &mut self.g,
+            GeneType::C(Some(Constant::M)) => &mut self.m,
+            GeneType::C(Some(Constant::O)) => &mut self.o,
+            GeneType::C(Some(Constant::T)) => &mut self.t,
         };
 
         match db.binary_search_by_key(&germline.name, |g| g.name.clone()) {
@@ -141,36 +155,64 @@ impl Chain {
             self.variable.iter().map(|g| g.alleles.len()).sum::<usize>(),
             self.joining.len(),
             self.joining.iter().map(|g| g.alleles.len()).sum::<usize>(),
-            self.constant.len(),
-            self.constant.iter().map(|g| g.alleles.len()).sum::<usize>(),
+            self.c.len()
+                + self.a.len()
+                + self.d.len()
+                + self.e.len()
+                + self.g.len()
+                + self.m.len()
+                + self.o.len()
+                + self.t.len(),
+            self.c.iter().map(|g| g.alleles.len()).sum::<usize>()
+                + self.a.iter().map(|g| g.alleles.len()).sum::<usize>()
+                + self.d.iter().map(|g| g.alleles.len()).sum::<usize>()
+                + self.e.iter().map(|g| g.alleles.len()).sum::<usize>()
+                + self.g.iter().map(|g| g.alleles.len()).sum::<usize>()
+                + self.m.iter().map(|g| g.alleles.len()).sum::<usize>()
+                + self.o.iter().map(|g| g.alleles.len()).sum::<usize>()
+                + self.t.iter().map(|g| g.alleles.len()).sum::<usize>(),
         )
     }
 }
 
 impl<'a> IntoIterator for &'a Chain {
-    type IntoIter = std::array::IntoIter<(GeneType, &'a [Germline]), 3>;
+    type IntoIter = std::array::IntoIter<(GeneType, &'a [Germline]), 10>;
     type Item = (GeneType, &'a [Germline]);
 
     fn into_iter(self) -> Self::IntoIter {
         [
             (GeneType::V, self.variable.as_slice()),
             (GeneType::J, self.joining.as_slice()),
-            (GeneType::C(None), self.constant.as_slice()),
+            (GeneType::C(None), self.c.as_slice()),
+            (GeneType::C(Some(Constant::A)), self.a.as_slice()),
+            (GeneType::C(Some(Constant::D)), self.d.as_slice()),
+            (GeneType::C(Some(Constant::E)), self.e.as_slice()),
+            (GeneType::C(Some(Constant::G)), self.g.as_slice()),
+            (GeneType::C(Some(Constant::M)), self.m.as_slice()),
+            (GeneType::C(Some(Constant::O)), self.o.as_slice()),
+            (GeneType::C(Some(Constant::T)), self.t.as_slice()),
         ]
         .into_iter()
     }
 }
 
 #[cfg(feature = "rayon")]
 impl<'a> IntoParallelIterator for &'a Chain {
-    type Iter = rayon::array::IntoIter<(GeneType, &'a [Germline]), 3>;
+    type Iter = rayon::array::IntoIter<(GeneType, &'a [Germline]), 10>;
     type Item = (GeneType, &'a [Germline]);
 
     fn into_par_iter(self) -> Self::Iter {
         [
             (GeneType::V, self.variable.as_slice()),
             (GeneType::J, self.joining.as_slice()),
-            (GeneType::C(None), self.constant.as_slice()),
+            (GeneType::C(None), self.c.as_slice()),
+            (GeneType::C(Some(Constant::A)), self.a.as_slice()),
+            (GeneType::C(Some(Constant::D)), self.d.as_slice()),
+            (GeneType::C(Some(Constant::E)), self.e.as_slice()),
+            (GeneType::C(Some(Constant::G)), self.g.as_slice()),
+            (GeneType::C(Some(Constant::M)), self.m.as_slice()),
+            (GeneType::C(Some(Constant::O)), self.o.as_slice()),
+            (GeneType::C(Some(Constant::T)), self.t.as_slice()),
         ]
         .into_par_iter()
     }