From a6da1637e7167d21c26dd0c5e13d0b8d49957e9a Mon Sep 17 00:00:00 2001
From: Sergey Bronnikov <estetus@gmail.com>
Date: Tue, 18 Jun 2024 13:57:26 +0300
Subject: [PATCH] molly: add a mix iterator

The patch add a function that takes a random mixture of a number
generators and chooses between them uniformly.
---
 CHANGELOG.md   |  1 +
 molly/gen.lua  | 53 ++++++++++++++++++++++++++++++++++++++++++--------
 test/tests.lua |  9 ++++++++-
 3 files changed, 54 insertions(+), 9 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f839a05..dda991b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- A `mix()` iterator.
 - A CAS-register generator.
 
 ### Changed
diff --git a/molly/gen.lua b/molly/gen.lua
index 3144c14..943ec8d 100644
--- a/molly/gen.lua
+++ b/molly/gen.lua
@@ -387,19 +387,56 @@ local cycle_times = function()
 end
 methods.cycle_times = cycle_times
 
---- (TODO) A random mixture of several generators. Takes a collection of generators
--- and chooses between them uniformly.
+local mix_gen
+
+mix_gen = function(_, state)
+    assert(type(state) == 'table')
+    local len = table.getn(state)
+    if len == 0 then
+        return nil, nil
+    end
+    local nth = math.random(len)
+    local it = state[nth]
+    local gen1, param1, state1 = unwrap(it)
+    local state2, value = gen1(param1, state1)
+    if value == nil then
+        table.remove(state, nth)
+        return mix_gen(nil, state)
+    end
+    state[nth] = fun.wrap(gen1, param1, state2)
+    return state, value
+end
+
+--- A random mixture of a number generators. Takes a collection of
+-- generators and chooses between them uniformly.
 --
--- To be precise, a mix behaves like a sequence of one-time, randomly selected
--- generators from the given collection. This is efficient and prevents
--- multiple generators from competing for the next slot, making it hard to
--- control the mixture of operations.
+-- @usage
+-- > molly.gen.range(1, 5):mix(molly.gen.range(5, 10)):totable()
+-- ---
+-- - - 1
+--   - 5
+--   - 2
+--   - 3
+--   - 6
+--   - 7
+--   - 4
+--   - 8
+--   - 9
+--   - 5
+--   - 10
 --
 -- @param ... - an iterators
 -- @return an iterator
 -- @function mix
-local mix = function()
-    -- TODO
+local function mix(...)
+    local params = {...}
+    local state = {}
+    for _, it in ipairs(params) do
+        if tostring(it) == '<generator>' then
+            table.insert(state, it)
+        end
+    end
+    return fun.wrap(mix_gen, nil, state)
 end
 methods.mix = mix
 exports.mix = mix
diff --git a/test/tests.lua b/test/tests.lua
index 176082a..b9cd18c 100644
--- a/test/tests.lua
+++ b/test/tests.lua
@@ -27,7 +27,7 @@ local utils = molly.utils
 local seed = os.time()
 math.randomseed(seed)
 
-test:plan(11)
+test:plan(12)
 
 test:test('clock', function(test)
     test:plan(5)
@@ -219,6 +219,13 @@ test:test('tests.list_append_gen', function(test)
     test:is(type(op_val) == 'number' or op_val == json.NULL, true, "tests.list_append_gen(): op value")
 end)
 
+test:test('gen.mix', function(test)
+    test:plan(1)
+
+    local tbl = molly.gen.range(1, 5):mix(molly.gen.range(5, 10)):totable()
+    test:is(#tbl, 11, 'length of items generated by gen.mix')
+end)
+
 test:test('runner', function(test)
     test:plan(4)