Skip to main content

minusone/ps/
array.rs

1use crate::error::{Error, MinusOneError, MinusOneErrorKind, MinusOneResult};
2use crate::ps::Powershell;
3use crate::ps::Powershell::{Array, Raw};
4use crate::ps::Value::Num;
5use crate::rule::RuleMut;
6use crate::tree::{ControlFlow, NodeMut};
7
8/// Parse array literal
9///
10/// It will parse 1,2,3,"5" as an array for powershell
11///
12/// # Example
13/// ```
14/// use minusone::ps::build_powershell_tree;
15/// use minusone::ps::forward::Forward;
16/// use minusone::ps::integer::ParseInt;
17/// use minusone::ps::linter::Linter;
18/// use minusone::ps::string::ParseString;
19/// use minusone::ps::access::AccessString;
20/// use minusone::ps::join::JoinOperator;
21/// use minusone::ps::array::ParseArrayLiteral;
22///
23/// let mut tree = build_powershell_tree("-join ('a','b','c')").unwrap();
24/// tree.apply_mut(&mut (
25///     ParseInt::default(),
26///     Forward::default(),
27///     ParseString::default(),
28///     ParseArrayLiteral::default(),
29///     JoinOperator::default()
30///     )
31/// ).unwrap();
32///
33/// let mut ps_litter_view = Linter::new();
34/// tree.apply(&mut ps_litter_view).unwrap();
35///
36/// assert_eq!(ps_litter_view.output, "\"abc\"");
37/// ```
38#[derive(Default)]
39pub struct ParseArrayLiteral;
40
41impl<'a> RuleMut<'a> for ParseArrayLiteral {
42    type Language = Powershell;
43
44    fn enter(
45        &mut self,
46        _node: &mut NodeMut<'a, Self::Language>,
47        _flow: ControlFlow,
48    ) -> MinusOneResult<()> {
49        Ok(())
50    }
51
52    fn leave(
53        &mut self,
54        node: &mut NodeMut<'a, Self::Language>,
55        _flow: ControlFlow,
56    ) -> MinusOneResult<()> {
57        let view = node.view();
58        if view.kind() == "array_literal_expression" && view.child_count() > 1 {
59            let mut range = vec![];
60            for child in view.iter() {
61                if let Some(Raw(value)) = child.data() {
62                    range.push(value.clone());
63                } else if child.kind() != "," {
64                    return Ok(());
65                }
66            }
67            node.set(Array(range));
68        }
69        Ok(())
70    }
71}
72
73/// This rule will generate
74/// a range value from operator ..
75///
76/// # Example
77/// ```
78/// extern crate tree_sitter;
79///
80/// use minusone::ps::build_powershell_tree;
81/// use minusone::ps::forward::Forward;
82/// use minusone::ps::integer::ParseInt;
83/// use minusone::ps::linter::Linter;
84/// use minusone::ps::string::ParseString;
85/// use minusone::ps::access::AccessString;
86/// use minusone::ps::join::JoinOperator;
87/// use minusone::ps::array::ParseRange;
88///
89/// let mut tree = build_powershell_tree("-join \"abc\"[0..2]").unwrap();
90/// tree.apply_mut(&mut (
91///     ParseInt::default(),
92///     Forward::default(),
93///     ParseRange::default(),
94///     ParseString::default(),
95///     JoinOperator::default(),
96///     AccessString::default()
97///     )
98/// ).unwrap();
99///
100/// let mut ps_litter_view = Linter::new();
101/// tree.apply(&mut ps_litter_view).unwrap();
102///
103/// assert_eq!(ps_litter_view.output, "\"abc\"");
104/// ```
105#[derive(Default)]
106pub struct ParseRange;
107
108impl<'a> RuleMut<'a> for ParseRange {
109    type Language = Powershell;
110
111    fn enter(
112        &mut self,
113        _node: &mut NodeMut<'a, Self::Language>,
114        _flow: ControlFlow,
115    ) -> MinusOneResult<()> {
116        Ok(())
117    }
118
119    fn leave(
120        &mut self,
121        node: &mut NodeMut<'a, Self::Language>,
122        _flow: ControlFlow,
123    ) -> MinusOneResult<()> {
124        let view = node.view();
125        if view.kind() == "range_expression" {
126            if let (Some(left_node), Some(right_node)) = (view.child(0), view.child(2)) {
127                if let (Some(Raw(left_value)), Some(Raw(right_value))) =
128                    (left_node.data(), right_node.data())
129                {
130                    if let (Some(from), Some(to)) =
131                        (left_value.clone().to_i64(), right_value.clone().to_i64())
132                    {
133                        let mut result = Vec::new();
134
135                        let mut index = from;
136                        let end = if from <= to { to + 1 } else { to - 1 };
137
138                        while index != end {
139                            result.push(Num(index));
140                            if from <= to {
141                                index += 1
142                            } else {
143                                index -= 1
144                            }
145                        }
146
147                        node.set(Array(result));
148                    }
149                }
150            }
151        }
152        Ok(())
153    }
154}
155
156/// This rule will handle array decared using @ operator
157///
158/// @(1, 2; 2+1)
159///
160/// # Example
161/// ```
162/// extern crate tree_sitter;
163/// extern crate tree_sitter_powershell;
164///
165/// use minusone::ps::build_powershell_tree;
166/// use minusone::ps::forward::Forward;
167/// use minusone::ps::integer::{ParseInt, AddInt};
168/// use minusone::ps::linter::Linter;
169/// use minusone::ps::string::ParseString;
170/// use minusone::ps::access::AccessString;
171/// use minusone::ps::join::JoinOperator;
172/// use minusone::ps::array::{ComputeArrayExpr, ParseArrayLiteral};
173///
174/// let mut tree = build_powershell_tree("-join \"abc\"[@(0, 1; 1+1)]").unwrap();
175/// tree.apply_mut(&mut (
176///     ParseInt::default(),
177///     AddInt::default(),
178///     Forward::default(),
179///     ComputeArrayExpr::default(),
180///     ParseString::default(),
181///     JoinOperator::default(),
182///     AccessString::default(),
183///     ParseArrayLiteral::default()
184///     )
185/// ).unwrap();
186///
187/// let mut ps_litter_view = Linter::new();
188/// tree.apply(&mut ps_litter_view).unwrap();
189///
190/// assert_eq!(ps_litter_view.output, "\"abc\"");
191/// ```
192#[derive(Default)]
193pub struct ComputeArrayExpr;
194
195impl<'a> RuleMut<'a> for ComputeArrayExpr {
196    type Language = Powershell;
197
198    fn enter(
199        &mut self,
200        _node: &mut NodeMut<'a, Self::Language>,
201        _flow: ControlFlow,
202    ) -> MinusOneResult<()> {
203        Ok(())
204    }
205
206    fn leave(
207        &mut self,
208        node: &mut NodeMut<'a, Self::Language>,
209        _flow: ControlFlow,
210    ) -> MinusOneResult<()> {
211        let view = node.view();
212        if view.kind() == "array_expression" {
213            if let Some(statement_list) = view.named_child("statements") {
214                let mut result = Vec::new();
215                for statement in statement_list.iter() {
216                    if statement.kind() == "empty_statement" {
217                        continue;
218                    }
219                    match statement.data() {
220                        Some(Array(values)) => {
221                            result.extend(values.clone());
222                        }
223                        Some(Raw(value)) => {
224                            result.push(value.clone());
225                        }
226                        _ => {
227                            // stop inferring
228                            return Ok(());
229                        }
230                    }
231                }
232                node.reduce(Array(result));
233            }
234        }
235        Ok(())
236    }
237}
238
239/// This rule will array concat using + operator
240///
241/// "foo"[0,1] + 'x' => ['f', 'o', 'x']
242///
243/// # Example
244/// ```
245/// extern crate tree_sitter;
246/// extern crate tree_sitter_powershell;
247///
248/// use minusone::ps::build_powershell_tree;
249/// use minusone::ps::forward::Forward;
250/// use minusone::ps::integer::{ParseInt, AddInt};
251/// use minusone::ps::linter::Linter;
252/// use minusone::ps::string::ParseString;
253/// use minusone::ps::access::AccessString;
254/// use minusone::ps::join::JoinOperator;
255/// use minusone::ps::array::{ComputeArrayExpr, ParseArrayLiteral, AddArray};
256///
257/// let mut tree = build_powershell_tree("-join ('foo'[0,1] + 'x')").unwrap();
258/// tree.apply_mut(&mut (
259///     ParseInt::default(),
260///     Forward::default(),
261///     ComputeArrayExpr::default(),
262///     ParseString::default(),
263///     JoinOperator::default(),
264///     AccessString::default(),
265///     ParseArrayLiteral::default(),
266///     AddArray::default()
267///     )
268/// ).unwrap();
269///
270/// let mut ps_litter_view = Linter::new();
271/// tree.apply(&mut ps_litter_view).unwrap();
272///
273/// assert_eq!(ps_litter_view.output, "\"fox\"");
274/// ```
275#[derive(Default)]
276pub struct AddArray;
277
278impl<'a> RuleMut<'a> for AddArray {
279    type Language = Powershell;
280
281    fn enter(
282        &mut self,
283        _node: &mut NodeMut<'a, Self::Language>,
284        _flow: ControlFlow,
285    ) -> MinusOneResult<()> {
286        Ok(())
287    }
288
289    fn leave(
290        &mut self,
291        node: &mut NodeMut<'a, Self::Language>,
292        _flow: ControlFlow,
293    ) -> MinusOneResult<()> {
294        let node_view = node.view();
295        if node_view.kind() == "additive_expression"
296            || node_view.kind() == "additive_argument_expression"
297        {
298            if let (Some(left_op), Some(operator), Some(right_op)) =
299                (node_view.child(0), node_view.child(1), node_view.child(2))
300            {
301                match (left_op.data(), operator.text()?, right_op.data()) {
302                    (Some(Array(array)), "+", Some(Raw(v))) => {
303                        let mut new_array = array.clone();
304                        new_array.push(v.clone());
305                        node.reduce(Array(new_array));
306                    }
307                    // add array between both
308                    (Some(Array(right_array)), "+", Some(Array(left_array))) => {
309                        node.reduce(Array([right_array.clone(), left_array.clone()].concat()));
310                    }
311                    _ => {}
312                }
313            }
314        }
315        Ok(())
316    }
317}
318
319/// This rule will generate an array resulting from a `New-Object` command invocation.
320///
321/// # Example
322/// ```
323/// extern crate tree_sitter;
324///
325/// use minusone::ps::build_powershell_tree;
326/// use minusone::ps::forward::Forward;
327/// use minusone::ps::integer::ParseInt;
328/// use minusone::ps::array::NewObjectArray;
329/// use minusone::ps::Value::Num;
330/// use minusone::ps::Powershell::Array;
331///
332/// let mut tree = build_powershell_tree("New-Object byte[] 16").unwrap();
333/// tree.apply_mut(&mut (
334///     ParseInt::default(),
335///     Forward::default(),
336///     NewObjectArray::default(),
337///     )
338/// ).unwrap();
339///
340/// assert_eq!(
341///     tree.root()
342///     .unwrap()
343///     .child(0)
344///     .unwrap()
345///     .child(0)
346///     .unwrap()
347///     .data(),
348///     Some(&Array(vec![Num(0); 16]))
349/// );
350/// ```
351#[derive(Default, Debug, Clone)]
352pub struct NewObjectArray {
353    max_size: Option<usize>,
354}
355
356impl NewObjectArray {
357    /// Returns a new [`NewObjectArray`] with no maximum capacity. Untrusted script input can cause
358    /// unbounded vec allocation.
359    pub fn new() -> Self {
360        Self { max_size: None }
361    }
362
363    /// Returns a [`NewObjectArray`] with a maximum capacity. Will infer array data no larger than
364    /// this upper bound. Invocations exceeding this value will cause an error to be returned
365    /// from [`RuleMut::leave`].
366    pub fn with_max_capacity(max_size: usize) -> Self {
367        Self {
368            max_size: Some(max_size),
369        }
370    }
371}
372
373impl<'a> RuleMut<'a> for NewObjectArray {
374    type Language = Powershell;
375
376    fn enter(
377        &mut self,
378        _node: &mut NodeMut<'a, Self::Language>,
379        _flow: ControlFlow,
380    ) -> MinusOneResult<()> {
381        Ok(())
382    }
383
384    fn leave(
385        &mut self,
386        node: &mut NodeMut<'a, Self::Language>,
387        _flow: ControlFlow,
388    ) -> MinusOneResult<()> {
389        let view = node.view();
390
391        if view.kind() != "command" || view.child_count() != 2 {
392            return Ok(());
393        }
394
395        if let (Some(command_name), Some(command_elements)) = (
396            view.named_child("command_name"),
397            view.named_child("command_elements"),
398        ) {
399            if command_name
400                .text()
401                .is_ok_and(|name| name.to_lowercase() == "new-object")
402                && command_elements.child_count() == 4
403                && command_elements
404                    .child(0)
405                    .is_some_and(|c| c.kind() == "command_argument_sep")
406                && command_elements.child(1).is_some_and(|c1| {
407                    c1.kind() == "generic_token"
408                        && matches!(c1.text(), Ok(t) if ["byte[]", "int[]"].iter().any(|typ| **typ == t.to_lowercase()))
409                })
410                && command_elements
411                    .child(2)
412                    .is_some_and(|c2| c2.kind() == "command_argument_sep")
413                && command_elements
414                    .child(3)
415                    .filter(|c| c.data().is_some())
416                    .is_some()
417            {
418                if let Some(Raw(Num(size))) =
419                    command_elements.child(3).as_ref().and_then(|c| c.data())
420                {
421                    if matches!(self.max_size, Some(max_size) if (*size) as usize > max_size) {
422                        return Err(Error::MinusOneError(
423                            MinusOneError::new(
424                                MinusOneErrorKind::InvalidProgram, 
425                                &format!("Array of length {} exceeds maximum array length", size)
426                            )
427                        ));
428                    }
429                    node.set(Array(vec![
430                        Num(0);
431                        std::cmp::max(
432                            (*size) as usize,
433                            self.max_size.unwrap_or_default()
434                        )
435                    ]));
436                }
437            }
438        }
439        Ok(())
440    }
441}
442
443#[cfg(test)]
444mod test {
445    use crate::ps::access::AccessString;
446    use crate::ps::array::{
447        AddArray, ComputeArrayExpr, NewObjectArray, ParseArrayLiteral, ParseRange,
448    };
449    use crate::ps::build_powershell_tree;
450    use crate::ps::forward::Forward;
451    use crate::ps::integer::{AddInt, ParseInt};
452    use crate::ps::string::ParseString;
453    use crate::ps::Powershell::Array;
454    use crate::ps::Value::{Num, Str};
455
456    #[test]
457    fn test_init_num_array() {
458        let mut tree = build_powershell_tree("@(1,2,3)").unwrap();
459        tree.apply_mut(&mut (
460            ParseInt::default(),
461            Forward::default(),
462            ComputeArrayExpr::default(),
463            ParseArrayLiteral::default(),
464        ))
465        .unwrap();
466
467        assert_eq!(
468            *tree
469                .root()
470                .unwrap()
471                .child(0)
472                .unwrap()
473                .child(0)
474                .unwrap()
475                .data()
476                .expect("Inferred type"),
477            Array(vec![Num(1), Num(2), Num(3)])
478        );
479    }
480
481    #[test]
482    fn test_init_mix_array() {
483        let mut tree = build_powershell_tree("@(1,2,'3')").unwrap();
484        tree.apply_mut(&mut (
485            ParseInt::default(),
486            ParseString::default(),
487            Forward::default(),
488            ComputeArrayExpr::default(),
489            ParseArrayLiteral::default(),
490        ))
491        .unwrap();
492
493        assert_eq!(
494            *tree
495                .root()
496                .unwrap()
497                .child(0)
498                .unwrap()
499                .child(0)
500                .unwrap()
501                .data()
502                .expect("Inferred type"),
503            Array(vec![Num(1), Num(2), Str("3".to_string())])
504        );
505    }
506
507    #[test]
508    fn test_init_str_array() {
509        let mut tree = build_powershell_tree("@('a','b','c')").unwrap();
510        tree.apply_mut(&mut (
511            ParseString::default(),
512            Forward::default(),
513            ComputeArrayExpr::default(),
514            ParseArrayLiteral::default(),
515        ))
516        .unwrap();
517
518        assert_eq!(
519            *tree
520                .root()
521                .unwrap()
522                .child(0)
523                .unwrap()
524                .child(0)
525                .unwrap()
526                .data()
527                .expect("Inferred type"),
528            Array(vec![
529                Str("a".to_string()),
530                Str("b".to_string()),
531                Str("c".to_string())
532            ])
533        );
534    }
535
536    #[test]
537    fn test_init_int_array_without_at() {
538        let mut tree = build_powershell_tree("1,2,3").unwrap();
539        tree.apply_mut(&mut (
540            ParseInt::default(),
541            Forward::default(),
542            ComputeArrayExpr::default(),
543            ParseArrayLiteral::default(),
544        ))
545        .unwrap();
546
547        assert_eq!(
548            *tree
549                .root()
550                .unwrap()
551                .child(0)
552                .unwrap()
553                .child(0)
554                .unwrap()
555                .data()
556                .expect("Inferred type"),
557            Array(vec![Num(1), Num(2), Num(3)])
558        );
559    }
560
561    #[test]
562    fn test_init_array_with_multi_statement() {
563        let mut tree = build_powershell_tree("@(1,2,3; 4 + 6)").unwrap();
564        tree.apply_mut(&mut (
565            ParseInt::default(),
566            AddInt::default(),
567            Forward::default(),
568            ComputeArrayExpr::default(),
569            ParseArrayLiteral::default(),
570        ))
571        .unwrap();
572
573        assert_eq!(
574            *tree
575                .root()
576                .unwrap()
577                .child(0)
578                .unwrap()
579                .child(0)
580                .unwrap()
581                .data()
582                .expect("Inferred type"),
583            Array(vec![Num(1), Num(2), Num(3), Num(10)])
584        );
585    }
586
587    #[test]
588    fn test_concat_array() {
589        let mut tree = build_powershell_tree("'foo'[0,1] + 'x'").unwrap();
590        tree.apply_mut(&mut (
591            ParseInt::default(),
592            AddArray::default(),
593            Forward::default(),
594            ComputeArrayExpr::default(),
595            ParseArrayLiteral::default(),
596            AccessString::default(),
597            ParseString::default(),
598        ))
599        .unwrap();
600
601        assert_eq!(
602            *tree
603                .root()
604                .unwrap()
605                .child(0)
606                .unwrap()
607                .child(0)
608                .unwrap()
609                .data()
610                .expect("Inferred type"),
611            Array(vec![
612                Str("f".to_string()),
613                Str("o".to_string()),
614                Str("x".to_string())
615            ])
616        );
617    }
618
619    #[test]
620    fn test_negative_range() {
621        let mut tree = build_powershell_tree("-1..-3").unwrap();
622        tree.apply_mut(&mut (
623            ParseInt::default(),
624            Forward::default(),
625            ComputeArrayExpr::default(),
626            ParseArrayLiteral::default(),
627            AccessString::default(),
628            ParseString::default(),
629            ParseRange::default(),
630        ))
631        .unwrap();
632
633        assert_eq!(
634            *tree
635                .root()
636                .unwrap()
637                .child(0)
638                .unwrap()
639                .child(0)
640                .unwrap()
641                .data()
642                .expect("Inferred type"),
643            Array(vec![Num(-1), Num(-2), Num(-3)])
644        );
645    }
646
647    #[test]
648    fn test_new_object_array() {
649        let mut tree = build_powershell_tree("New-Object byte[] 16").unwrap();
650        tree.apply_mut(&mut (ParseInt::default(), Forward::default()))
651            .unwrap();
652        tree.apply_mut(&mut NewObjectArray::default()).unwrap();
653
654        assert_eq!(
655            *tree
656                .root()
657                .unwrap()
658                .child(0)
659                .unwrap()
660                .child(0)
661                .unwrap()
662                .child(0)
663                .unwrap()
664                .data()
665                .expect("Inferred data"),
666            Array(vec![Num(0); 16]),
667        );
668    }
669}