From 0787bdfba3160222d02fb6b1e3e13f82576a7694 Mon Sep 17 00:00:00 2001
From: Akosh Farkash <aakoshh@gmail.com>
Date: Mon, 12 May 2025 19:22:12 +0100
Subject: [PATCH] Do not use zero length types in the main input output

---
 tooling/ast_fuzzer/src/program/func.rs  |  2 +-
 tooling/ast_fuzzer/src/program/mod.rs   | 32 +++++++++++++++++++++----
 tooling/ast_fuzzer/src/program/types.rs | 12 ++++++++++
 3 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/tooling/ast_fuzzer/src/program/func.rs b/tooling/ast_fuzzer/src/program/func.rs
index b7e064d4839..defa6975af9 100644
--- a/tooling/ast_fuzzer/src/program/func.rs
+++ b/tooling/ast_fuzzer/src/program/func.rs
@@ -715,7 +715,7 @@ impl<'a> FunctionContext<'a> {
         // Generate a type or choose an existing one.
         let max_depth = self.max_depth();
         let comptime_friendly = self.is_comptime_friendly();
-        let typ = self.ctx.gen_type(u, max_depth, false, true, comptime_friendly)?;
+        let typ = self.ctx.gen_type(u, max_depth, false, false, true, comptime_friendly)?;
         let expr = self.gen_expr(u, &typ, max_depth, Flags::TOP)?;
         let mutable = bool::arbitrary(u)?;
         Ok(self.let_var(mutable, typ, expr, true))
diff --git a/tooling/ast_fuzzer/src/program/mod.rs b/tooling/ast_fuzzer/src/program/mod.rs
index 34aa814c6df..bf5f388ac67 100644
--- a/tooling/ast_fuzzer/src/program/mod.rs
+++ b/tooling/ast_fuzzer/src/program/mod.rs
@@ -133,8 +133,14 @@ impl Context {
         u: &mut Unstructured,
         i: usize,
     ) -> arbitrary::Result<(Name, Type, Expression)> {
-        let typ =
-            self.gen_type(u, self.config.max_depth, true, false, self.config.comptime_friendly)?;
+        let typ = self.gen_type(
+            u,
+            self.config.max_depth,
+            true,
+            false,
+            false,
+            self.config.comptime_friendly,
+        )?;
         // By the time we get to the monomorphized AST the compiler will have already turned
         // complex global expressions into literals.
         let val = expr::gen_literal(u, &typ)?;
@@ -172,6 +178,7 @@ impl Context {
                 u,
                 self.config.max_depth,
                 false,
+                is_main,
                 false,
                 self.config.comptime_friendly,
             )?;
@@ -189,8 +196,15 @@ impl Context {
             params.push((id, is_mutable, name, typ, visibility));
         }
 
-        let return_type =
-            self.gen_type(u, self.config.max_depth, false, false, self.config.comptime_friendly)?;
+        let return_type = self.gen_type(
+            u,
+            self.config.max_depth,
+            false,
+            is_main,
+            false,
+            self.config.comptime_friendly,
+        )?;
+
         let return_visibility = if is_main {
             if types::is_unit(&return_type) {
                 Visibility::Private
@@ -321,6 +335,7 @@ impl Context {
         u: &mut Unstructured,
         max_depth: usize,
         is_global: bool,
+        is_main: bool,
         is_frontend_friendly: bool,
         is_comptime_friendly: bool,
     ) -> arbitrary::Result<Type> {
@@ -330,6 +345,7 @@ impl Context {
                 .types
                 .iter()
                 .filter(|typ| !is_global || types::can_be_global(typ))
+                .filter(|typ| !is_main || !types::can_be_main(typ))
                 .filter(|typ| types::type_depth(typ) <= max_depth)
                 .filter(|typ| !is_frontend_friendly || !self.should_avoid_literals(typ))
                 .collect::<Vec<_>>();
@@ -378,6 +394,7 @@ impl Context {
                                 u,
                                 max_depth - 1,
                                 is_global,
+                                is_main,
                                 is_frontend_friendly,
                                 is_comptime_friendly,
                             )
@@ -392,6 +409,7 @@ impl Context {
                         u,
                         max_depth - 1,
                         is_global,
+                        is_main,
                         is_frontend_friendly,
                         is_comptime_friendly,
                     )?;
@@ -399,7 +417,11 @@ impl Context {
                 }
                 _ => unreachable!("unexpected arbitrary type index"),
             };
-            if !is_global || types::can_be_global(&typ) {
+            // Looping is kinda dangerous, we could get stuck if we run out of randomness,
+            // so we have to make sure the first type on the list is acceptable.
+            if is_global && !types::can_be_global(&typ) || is_main && !types::can_be_main(&typ) {
+                continue;
+            } else {
                 break;
             }
         }
diff --git a/tooling/ast_fuzzer/src/program/types.rs b/tooling/ast_fuzzer/src/program/types.rs
index 229a06c0161..eef9afc3ee7 100644
--- a/tooling/ast_fuzzer/src/program/types.rs
+++ b/tooling/ast_fuzzer/src/program/types.rs
@@ -35,6 +35,18 @@ pub(crate) fn can_be_global(typ: &Type) -> bool {
     )
 }
 
+/// Check if a type can be used in the `main` function.
+///
+/// We decided we will avoid 0 length arrays in the main inputs and outputs, because we don't generate
+/// witnesses for them anyway, and they are tricky to handle consistently when they can be regular inputs
+/// as well as part of the databus. They are not expected in real programs as they don't do anything useful.
+pub(crate) fn can_be_main(typ: &Type) -> bool {
+    match typ {
+        Type::Array(size, _) | Type::String(size) => *size > 0,
+        _ => true,
+    }
+}
+
 /// Collect all the sub-types produced by a type.
 ///
 /// It's like a _power set_ of the type.