From 0f521821d55f000acb5e63e9251274fc478ff456 Mon Sep 17 00:00:00 2001
From: Yi2255 <cjy096010@gmail.com>
Date: Fri, 15 Nov 2024 17:43:52 +0800
Subject: [PATCH] Fix non-iteratrable input of yieldEach (#463)

---
 .../CodeGen/CodeGeneratorWeights.swift        |  3 ++-
 Sources/Fuzzilli/CodeGen/CodeGenerators.swift | 24 +++++++++++--------
 2 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/Sources/Fuzzilli/CodeGen/CodeGeneratorWeights.swift b/Sources/Fuzzilli/CodeGen/CodeGeneratorWeights.swift
index 9ce5be34..bda9a0d7 100644
--- a/Sources/Fuzzilli/CodeGen/CodeGeneratorWeights.swift
+++ b/Sources/Fuzzilli/CodeGen/CodeGeneratorWeights.swift
@@ -90,7 +90,8 @@ public let codeGeneratorWeights = [
     "AsyncGeneratorFunctionGenerator":          1,
     "ConstructorGenerator":                     5,
     "SubroutineReturnGenerator":                3,
-    "YieldGenerator":                           2,
+    "YieldGenerator":                           5,
+    "YieldEachGenerator":                       5,
     "AwaitGenerator":                           2,
     "PropertyRetrievalGenerator":               20,
     "PropertyAssignmentGenerator":              20,
diff --git a/Sources/Fuzzilli/CodeGen/CodeGenerators.swift b/Sources/Fuzzilli/CodeGen/CodeGenerators.swift
index f8a9cd5f..616106d9 100644
--- a/Sources/Fuzzilli/CodeGen/CodeGenerators.swift
+++ b/Sources/Fuzzilli/CodeGen/CodeGenerators.swift
@@ -775,7 +775,9 @@ public let CodeGenerators: [CodeGenerator] = [
             if probability(0.5) {
                 b.yield(b.randomVariable())
             } else {
-                b.yieldEach(b.randomVariable())
+                let randomVariables = b.randomVariables(n: Int.random(in: 1...5))
+                let array = b.createArray(with: randomVariables)
+                b.yieldEach(array)
             }
             b.doReturn(b.randomVariable())
         }
@@ -807,7 +809,9 @@ public let CodeGenerators: [CodeGenerator] = [
             if probability(0.5) {
                 b.yield(b.randomVariable())
             } else {
-                b.yieldEach(b.randomVariable())
+                let randomVariables = b.randomVariables(n: Int.random(in: 1...5))
+                let array = b.createArray(with: randomVariables)
+                b.yieldEach(array)
             }
             b.doReturn(b.randomVariable())
         }
@@ -1077,18 +1081,18 @@ public let CodeGenerators: [CodeGenerator] = [
 
     CodeGenerator("YieldGenerator", inContext: .generatorFunction, inputs: .one) { b, val in
         assert(b.context.contains(.generatorFunction))
-        if probability(0.5) {
-            if probability(0.9) {
-                b.yield(val)
-            } else {
-                b.yield()
-            }
+        if probability(0.9) {
+            b.yield(val)
         } else {
-            // TODO only do this when the value is iterable?
-            b.yieldEach(val)
+            b.yield()
         }
     },
 
+    CodeGenerator("YieldEachGenerator", inContext: .generatorFunction, inputs: .required(.iterable)) { b, val in
+        assert(b.context.contains(.generatorFunction))
+        b.yieldEach(val)
+    },
+
     CodeGenerator("AwaitGenerator", inContext: .asyncFunction, inputs: .one) { b, val in
         assert(b.context.contains(.asyncFunction))
         b.await(val)