conspire/geometry/ntree/leaves/
mod.rs

1use crate::geometry::ntree::{
2    Orthotree,
3    node::{Kind, Node},
4    subdivide::insert_bit,
5};
6use std::array::from_fn;
7
8impl<const D: usize, const L: usize, const M: usize, const N: usize, T, U>
9    Orthotree<D, L, M, N, T, U>
10where
11    T: Copy + Into<usize>,
12    U: Copy + Into<usize>,
13{
14    pub fn all_leaves<'a>(&self, node: &'a Node<D, M, N, T, U>) -> Option<&'a [U; N]> {
15        match &node.kind {
16            Kind::Leaf => None,
17            Kind::Tree(orthants) => {
18                if orthants.iter().any(|&orthant| self[orthant].is_tree()) {
19                    None
20                } else {
21                    Some(orthants)
22                }
23            }
24        }
25    }
26    pub fn leaves(&self, node: &Node<D, M, N, T, U>) -> [Option<U>; N] {
27        match &node.kind {
28            Kind::Leaf => [None; N],
29            Kind::Tree(orthants) => from_fn(|i| {
30                if self[orthants[i]].is_leaf() {
31                    Some(orthants[i])
32                } else {
33                    None
34                }
35            }),
36        }
37    }
38    pub fn leaves_on_facet(&self, node: &Node<D, M, N, T, U>, facet: usize) -> [Option<U>; L] {
39        let (axis, side) = (facet >> 1, facet & 1);
40        match &node.kind {
41            Kind::Leaf => from_fn(|_| None),
42            Kind::Tree(orthants) => from_fn(|i| {
43                let orthant = orthants[insert_bit(i, axis, side)];
44                self[orthant].is_leaf().then_some(orthant)
45            }),
46        }
47    }
48    pub fn leaves_and_facets(
49        &self,
50        node: &Node<D, M, N, T, U>,
51    ) -> [Option<(U, [Option<U>; D])>; N] {
52        match &node.kind {
53            Kind::Leaf => from_fn(|_| None),
54            Kind::Tree(orthants) => from_fn(|i| {
55                let orthant = orthants[i];
56                if self[orthant].is_leaf() {
57                    let facets = &self[orthant].facets;
58                    let external: [Option<U>; D] = from_fn(|b| facets[2 * b + ((i >> b) & 1)]);
59                    Some((orthant, external))
60                } else {
61                    None
62                }
63            }),
64        }
65    }
66    pub fn orthants_leaves(&self, node: &Node<D, M, N, T, U>) -> [Option<[Option<U>; N]>; N] {
67        match &node.kind {
68            Kind::Leaf => from_fn(|_| None),
69            Kind::Tree(orthants) => from_fn(|i| match &self[orthants[i]].kind {
70                Kind::Leaf => None,
71                Kind::Tree(sub_orthants) => {
72                    let inner: [Option<U>; N] = from_fn(|j| {
73                        if self[sub_orthants[j]].is_leaf() {
74                            Some(sub_orthants[j])
75                        } else {
76                            None
77                        }
78                    });
79                    if inner.iter().any(|x| x.is_some()) {
80                        Some(inner)
81                    } else {
82                        None
83                    }
84                }
85            }),
86        }
87    }
88    pub fn orthants_leaves_on_facet(
89        &self,
90        node: &Node<D, M, N, T, U>,
91        face: usize,
92    ) -> [Option<[Option<U>; L]>; L] {
93        let (axis, side) = (face >> 1, face & 1);
94        match &node.kind {
95            Kind::Leaf => from_fn(|_| None),
96            Kind::Tree(orthants) => {
97                from_fn(|i| match &self[orthants[insert_bit(i, axis, side)]].kind {
98                    Kind::Leaf => None,
99                    Kind::Tree(sub_orthants) => {
100                        let inner: [Option<U>; L] = from_fn(|j| {
101                            let leaf = sub_orthants[insert_bit(j, axis, side)];
102                            self[leaf].is_leaf().then_some(leaf)
103                        });
104                        inner.iter().any(|x| x.is_some()).then_some(inner)
105                    }
106                })
107            }
108        }
109    }
110    pub fn orthants_all_leaves_on_facet(
111        &self,
112        node: &Node<D, M, N, T, U>,
113        face: usize,
114    ) -> Option<[[U; L]; L]> {
115        let orthants_leaves = self.orthants_leaves_on_facet(node, face);
116        orthants_leaves
117            .iter()
118            .all(|&orthant| orthant.is_some_and(|leaves| leaves.iter().all(Option::is_some)))
119            .then(|| from_fn(|i| from_fn(|j| orthants_leaves[i].unwrap()[j].unwrap())))
120    }
121}
conspire/geometry/ntree/leaves/mod.rs

conspire/geometry/ntree/leaves/
mod.rs
