From 28659d068705aded4405b18b844d50f4dc05e3d9 Mon Sep 17 00:00:00 2001
From: Yi2255 <cjy096010@gmail.com>
Date: Thu, 24 Oct 2024 00:33:13 +0000
Subject: [PATCH] Fix non-iteratrable input of yieldEach

---
 Sources/Fuzzilli/Base/ProgramBuilder.swift    |  9 +++++++++
 Sources/Fuzzilli/CodeGen/CodeGenerators.swift | 10 +++++-----
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/Sources/Fuzzilli/Base/ProgramBuilder.swift b/Sources/Fuzzilli/Base/ProgramBuilder.swift
index 3eb4cb968..04fb8fa88 100644
--- a/Sources/Fuzzilli/Base/ProgramBuilder.swift
+++ b/Sources/Fuzzilli/Base/ProgramBuilder.swift
@@ -2711,4 +2711,13 @@ public class ProgramBuilder {
             break
         }
     }
+
+    public func buildIteratorVariable(_ b: ProgramBuilder, _ it: Variable) -> Variable{
+        if(b.type(of: it).Is(.iterable)){
+            return it
+        }
+        let initialValues = (0..<Int.random(in: 1...5)).map({ _ in b.randomVariable() })
+        return b.createArray(with: initialValues)
+    }
+
 }
diff --git a/Sources/Fuzzilli/CodeGen/CodeGenerators.swift b/Sources/Fuzzilli/CodeGen/CodeGenerators.swift
index f8a9cd5f7..8934618a7 100644
--- a/Sources/Fuzzilli/CodeGen/CodeGenerators.swift
+++ b/Sources/Fuzzilli/CodeGen/CodeGenerators.swift
@@ -769,13 +769,13 @@ public let CodeGenerators: [CodeGenerator] = [
         // These are "typically" used as arguments, so we don't directly generate a call operation here.
     },
 
-    RecursiveCodeGenerator("GeneratorFunctionGenerator") { b in
+    RecursiveCodeGenerator("GeneratorFunctionGenerator", inputs: .preferred(.iterable)) { b, it in
         let f = b.buildGeneratorFunction(with: b.randomParameters(), isStrict: probability(0.1)) { _ in
             b.buildRecursive()
             if probability(0.5) {
                 b.yield(b.randomVariable())
             } else {
-                b.yieldEach(b.randomVariable())
+                b.yieldEach(b.buildIteratorVariable(b, it))
             }
             b.doReturn(b.randomVariable())
         }
@@ -800,14 +800,14 @@ public let CodeGenerators: [CodeGenerator] = [
         // These are "typically" used as arguments, so we don't directly generate a call operation here.
     },
 
-    RecursiveCodeGenerator("AsyncGeneratorFunctionGenerator") { b in
+    RecursiveCodeGenerator("AsyncGeneratorFunctionGenerator", inputs: .preferred(.iterable)) { b, it in
         let f = b.buildAsyncGeneratorFunction(with: b.randomParameters(), isStrict: probability(0.1)) { _ in
             b.buildRecursive()
             b.await(b.randomVariable())
             if probability(0.5) {
                 b.yield(b.randomVariable())
             } else {
-                b.yieldEach(b.randomVariable())
+                b.yieldEach(b.buildIteratorVariable(b, it))
             }
             b.doReturn(b.randomVariable())
         }
@@ -1085,7 +1085,7 @@ public let CodeGenerators: [CodeGenerator] = [
             }
         } else {
             // TODO only do this when the value is iterable?
-            b.yieldEach(val)
+            b.yieldEach(b.buildIteratorVariable(b, val))
         }
     },