1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

use crate::{Expression, Size, TensorExpression};
use std::collections::HashSet;

pub fn new_variable_tensor(id: String, sizes: Vec<Size>) -> Expression {
    Expression::Variable(id, sizes)
}

impl TensorExpression {
    pub fn variable_ids(&self) -> HashSet<&str> {
        match self {
            TensorExpression::KroneckerDeltas(_) => HashSet::new(),
            TensorExpression::DotProduct {
                terms,
                rank_combinations: _,
            } => terms.iter().map(|t| t.variable_ids()).flatten().collect(),
            TensorExpression::DirectProduct(terms) => {
                terms.iter().map(|t| t.variable_ids()).flatten().collect()
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashSet;

    use crate::{new_variable_tensor, size, MatrixExpression, Size};

    #[test]
    fn it_works() {
        let id = "x";
        let a = HashSet::from([id; 1]);
        let ea = new_variable_tensor((id).to_string(), vec![Size::Many, Size::Many, Size::Many]);
        println!("{:?}", ea);
    }
}