Support 2-argument form of ArrayQ with depth check

ArrayQ[expr, n] now checks if expr is a rectangular array of
depth n. E.g. ArrayQ[{{1,2},{3,4}}, 2] returns True.
Previously only the 1-argument form was supported.

Author: ad-si

src/evaluator/dispatch/arg_count.rs

@@ -65,7 +65,7 @@ pub fn get_arg_count_range(name: &str) -> Option<(usize, usize)> {
     "ArrayDepth" => Some((1, 1)),
     "ArrayFlatten" => Some((1, 1)),
     "ArrayPad" => Some((2, 3)),
-    "ArrayQ" => Some((1, 1)),
+    "ArrayQ" => Some((1, 3)),
     "ArrayResample" => Some((2, 2)),
     "ArrayReshape" => Some((2, 3)),
     "ArrayRules" => Some((1, 2)),

src/evaluator/dispatch/list_operations.rs

@@ -1458,8 +1458,27 @@ pub fn dispatch_list_operations(
     "TensorRank" if args.len() == 1 => {
       return Some(list_helpers_ast::tensor_rank_ast(&args[0]));
     }
-    "ArrayQ" if args.len() == 1 => {
-      return Some(list_helpers_ast::array_q_ast(&args[0]));
+    "ArrayQ" if !args.is_empty() && args.len() <= 3 => {
+      let is_array = match list_helpers_ast::array_q_ast(&args[0]) {
+        Ok(v) => v,
+        Err(e) => return Some(Err(e)),
+      };
+      if args.len() >= 2 {
+        // ArrayQ[expr, n] - check if expr is an array of depth n
+        if matches!(&is_array, Expr::Identifier(s) if s == "True")
+          && let Ok(dims) = list_helpers_ast::dimensions_ast(&[args[0].clone()])
+          && let Expr::List(d) = &dims
+          && let Some(n) = expr_to_i128(&args[1])
+        {
+          return Some(Ok(if d.len() == n as usize {
+            Expr::Identifier("True".to_string())
+          } else {
+            Expr::Identifier("False".to_string())
+          }));
+        }
+        return Some(Ok(is_array));
+      }
+      return Some(Ok(is_array));
     }
     "VectorQ" if args.len() == 1 => {
       return Some(list_helpers_ast::vector_q_ast(&args[0]));

tests/interpreter_tests/functions.rs

@@ -882,6 +882,21 @@ mod array_predicates {
     assert_eq!(interpret("ArrayQ[{{1, 2}, {3}}]").unwrap(), "False");
   }
 
+  #[test]
+  fn array_q_with_depth() {
+    assert_eq!(interpret("ArrayQ[{{1, 2}, {3, 4}}, 2]").unwrap(), "True");
+  }
+
+  #[test]
+  fn array_q_with_wrong_depth() {
+    assert_eq!(interpret("ArrayQ[{{1, 2}, {3, 4}}, 1]").unwrap(), "False");
+  }
+
+  #[test]
+  fn array_q_vector_depth_1() {
+    assert_eq!(interpret("ArrayQ[{1, 2, 3}, 1]").unwrap(), "True");
+  }
+
   #[test]
   fn matrix_q_true() {
     assert_eq!(interpret("MatrixQ[{{1, 2}, {3, 4}}]").unwrap(), "True");