Skip to main content

conspire/geometry/ntree/leaves/
mod.rs

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