Add shuffle and repeat options for the playback queue

6 files changed, 272 insertions, 43 deletions

diff --git a/src/audio/CMakeLists.txt b/src/audio/CMakeLists.txt
index 0f90334b..b219ab6e 100644
--- a/src/audio/CMakeLists.txt
+++ b/src/audio/CMakeLists.txt
@@ -9,6 +9,6 @@ idf_component_register(
   "audio_source.cpp"
   INCLUDE_DIRS "include"
   REQUIRES "codecs" "drivers" "cbor" "result" "tasks" "span" "memory" "tinyfsm"
-  "database" "system_fsm" "speexdsp")
+  "database" "system_fsm" "speexdsp" "millershuffle")
 
 target_compile_options(${COMPONENT_LIB} PRIVATE ${EXTRA_WARNINGS})
diff --git a/src/audio/include/track_queue.hpp b/src/audio/include/track_queue.hpp
index 4a1984c9..24b4fe48 100644
--- a/src/audio/include/track_queue.hpp
+++ b/src/audio/include/track_queue.hpp
@@ -19,6 +19,30 @@
 namespace audio {
 
 /*
+ * Utility that uses a Miller shuffle to yield well-distributed random indexes
+ * from within a range.
+ */
+class RandomIterator {
+ public:
+  RandomIterator(size_t size);
+
+  auto current() const -> size_t;
+
+  auto next() -> void;
+  auto prev() -> void;
+
+  // Note resizing has the side-effect of restarting iteration.
+  auto resize(size_t) -> void;
+  auto repeat(bool) -> void;
+
+ private:
+  size_t seed_;
+  size_t pos_;
+  size_t size_;
+  bool repeat_;
+};
+
+/*
  * Owns and manages a complete view of the playback queue. Includes the
  * currently playing track, a truncated list of previously played tracks, and
  * all future tracks that have been queued.
@@ -51,7 +75,7 @@ class TrackQueue {
   auto totalSize() const -> size_t;
 
   using Item = std::variant<database::TrackId, database::TrackIterator>;
-  auto insert(Item) -> void;
+  auto insert(Item, size_t index = 0) -> void;
   auto append(Item i) -> void;
 
   /*
@@ -68,6 +92,12 @@ class TrackQueue {
    */
   auto clear() -> void;
 
+  auto random(bool) -> void;
+  auto random() const -> bool;
+
+  auto repeat(bool) -> void;
+  auto repeat() const -> bool;
+
   // Cannot be copied or moved.
   TrackQueue(const TrackQueue&) = delete;
   TrackQueue& operator=(const TrackQueue&) = delete;
@@ -79,6 +109,9 @@ class TrackQueue {
 
   size_t pos_;
   std::pmr::vector<database::TrackId> tracks_;
+
+  std::optional<RandomIterator> shuffle_;
+  bool repeat_;
 };
 
 }  // namespace audio
diff --git a/src/audio/track_queue.cpp b/src/audio/track_queue.cpp
index c1187107..7e08e3a2 100644
--- a/src/audio/track_queue.cpp
+++ b/src/audio/track_queue.cpp
@@ -5,14 +5,18 @@
  */
 
 #include "track_queue.hpp"
-#include <stdint.h>
 
 #include <algorithm>
+#include <cstdint>
 #include <memory>
 #include <mutex>
 #include <optional>
+#include <shared_mutex>
 #include <variant>
 
+#include "MillerShuffle.h"
+#include "esp_random.h"
+
 #include "audio_events.hpp"
 #include "audio_fsm.hpp"
 #include "cppbor.h"
@@ -28,6 +32,42 @@ namespace audio {
 
 [[maybe_unused]] static constexpr char kTag[] = "tracks";
 
+RandomIterator::RandomIterator(size_t size)
+    : seed_(), pos_(0), size_(size), repeat_(false) {
+  esp_fill_random(&seed_, sizeof(seed_));
+}
+
+auto RandomIterator::current() const -> size_t {
+  if (pos_ < size_ || repeat_) {
+    return MillerShuffle(pos_, seed_, size_);
+  }
+  return size_;
+}
+
+auto RandomIterator::next() -> void {
+  // MillerShuffle behaves well with pos > size, returning different
+  // permutations each 'cycle'. We therefore don't need to worry about wrapping
+  // this value.
+  pos_++;
+}
+
+auto RandomIterator::prev() -> void {
+  if (pos_ > 0) {
+    pos_--;
+  }
+}
+
+auto RandomIterator::resize(size_t s) -> void {
+  size_ = s;
+  // Changing size will yield a different current position anyway, so reset pos
+  // to ensure we yield a full sweep of both new and old indexes.
+  pos_ = 0;
+}
+
+auto RandomIterator::repeat(bool r) -> void {
+  repeat_ = r;
+}
+
 auto notifyChanged(bool current_changed) -> void {
   QueueUpdate ev{.current_changed = current_changed};
   events::Ui().Dispatch(ev);
@@ -38,7 +78,9 @@ TrackQueue::TrackQueue(tasks::Worker& bg_worker)
     : mutex_(),
       bg_worker_(bg_worker),
       pos_(0),
-      tracks_(&memory::kSpiRamResource) {}
+      tracks_(&memory::kSpiRamResource),
+      shuffle_(),
+      repeat_(false) {}
 
 auto TrackQueue::current() const -> std::optional<database::TrackId> {
   const std::shared_lock<std::shared_mutex> lock(mutex_);
@@ -78,87 +120,191 @@ auto TrackQueue::totalSize() const -> size_t {
   return tracks_.size();
 }
 
-auto TrackQueue::insert(Item i) -> void {
-  bool current_changed = pos_ == tracks_.size();
+auto TrackQueue::insert(Item i, size_t index) -> void {
+  bool was_queue_empty;
+  bool current_changed;
+  {
+    const std::shared_lock<std::shared_mutex> lock(mutex_);
+    was_queue_empty = pos_ == tracks_.size();
+    current_changed = pos_ == was_queue_empty || index == pos_;
+  }
+
+  auto update_shuffler = [=, this]() {
+    if (shuffle_) {
+      shuffle_->resize(tracks_.size());
+      // If there wasn't anything already playing, then we should make sure we
+      // begin playback at a random point, instead of always starting with
+      // whatever was inserted first and *then* shuffling.
+      // We don't base this purely off of current_changed because we would like
+      // 'play this track now' (by inserting at the current pos) to work even
+      // when shuffling is enabled.
+      if (was_queue_empty) {
+        pos_ = shuffle_->current();
+      }
+    }
+  };
+
   if (std::holds_alternative<database::TrackId>(i)) {
-    const std::unique_lock<std::shared_mutex> lock(mutex_);
-    tracks_.push_back(std::get<database::TrackId>(i));
+    {
+      const std::unique_lock<std::shared_mutex> lock(mutex_);
+      if (index <= tracks_.size()) {
+        tracks_.insert(tracks_.begin() + index, std::get<database::TrackId>(i));
+        update_shuffler();
+      }
+    }
     notifyChanged(current_changed);
   } else if (std::holds_alternative<database::TrackIterator>(i)) {
+    // Iterators can be very large, and retrieving items from them often
+    // requires disk i/o. Handle them asynchronously so that inserting them
+    // doesn't block.
     bg_worker_.Dispatch<void>([=, this]() {
       database::TrackIterator it = std::get<database::TrackIterator>(i);
-      size_t working_pos = pos_;
+      size_t working_pos = index;
       while (true) {
         auto next = *it;
         if (!next) {
           break;
         }
-        const std::unique_lock<std::shared_mutex> lock(mutex_);
-        tracks_.insert(tracks_.begin() + working_pos, *next);
+        // Keep this critical section small so that we're not blocking methods
+        // like current().
+        {
+          const std::unique_lock<std::shared_mutex> lock(mutex_);
+          if (working_pos <= tracks_.size()) {
+            tracks_.insert(tracks_.begin() + working_pos, *next);
+          }
+        }
         working_pos++;
         it++;
       }
+      {
+        const std::unique_lock<std::shared_mutex> lock(mutex_);
+        update_shuffler();
+      }
       notifyChanged(current_changed);
     });
   }
 }
 
 auto TrackQueue::append(Item i) -> void {
-  bool current_changed = pos_ == tracks_.size();
-  if (std::holds_alternative<database::TrackId>(i)) {
-    const std::unique_lock<std::shared_mutex> lock(mutex_);
-    tracks_.push_back(std::get<database::TrackId>(i));
-    notifyChanged(current_changed);
-  } else if (std::holds_alternative<database::TrackIterator>(i)) {
-    bg_worker_.Dispatch<void>([=, this]() {
-      database::TrackIterator it = std::get<database::TrackIterator>(i);
-      while (true) {
-        auto next = *it;
-        if (!next) {
-          break;
-        }
-        const std::unique_lock<std::shared_mutex> lock(mutex_);
-        tracks_.push_back(*next);
-        it++;
-      }
-      notifyChanged(current_changed);
-    });
+  size_t end;
+  {
+    const std::shared_lock<std::shared_mutex> lock(mutex_);
+    end = tracks_.size();
   }
+  insert(i, end);
 }
 
 auto TrackQueue::next() -> void {
   const std::unique_lock<std::shared_mutex> lock(mutex_);
-  pos_ = std::min<size_t>(pos_ + 1, tracks_.size());
+  if (shuffle_) {
+    shuffle_->next();
+    pos_ = shuffle_->current();
+  } else {
+    pos_++;
+    if (pos_ >= tracks_.size() && repeat_) {
+      pos_ = 0;
+    }
+  }
 
   notifyChanged(true);
 }
 
 auto TrackQueue::previous() -> void {
   const std::unique_lock<std::shared_mutex> lock(mutex_);
-  if (pos_ > 0) {
-    pos_--;
+  if (shuffle_) {
+    shuffle_->prev();
+    pos_ = shuffle_->current();
+  } else {
+    if (pos_ == 0) {
+      if (repeat_) {
+        pos_ = tracks_.size() - 1;
+      }
+    } else {
+      pos_--;
+    }
   }
 
   notifyChanged(true);
 }
 
 auto TrackQueue::skipTo(database::TrackId id) -> void {
-  const std::unique_lock<std::shared_mutex> lock(mutex_);
-  for (size_t i = pos_; i < tracks_.size(); i++) {
-    if (tracks_[i] == id) {
-      pos_ = i;
+  // Defer this work to the background not because it's particularly
+  // long-running (although it could be), but because we want to ensure we only
+  // search for the given id after any previously pending iterator insertions
+  // have finished.
+  bg_worker_.Dispatch<void>([=, this]() {
+    bool found = false;
+    {
+      const std::unique_lock<std::shared_mutex> lock(mutex_);
+      for (size_t i = 0; i < tracks_.size(); i++) {
+        if (tracks_[i] == id) {
+          pos_ = i;
+          found = true;
+          break;
+        }
+      }
+    }
+    if (found) {
+      notifyChanged(true);
+    }
+  });
+}
+
+auto TrackQueue::clear() -> void {
+  {
+    const std::unique_lock<std::shared_mutex> lock(mutex_);
+    if (tracks_.empty()) {
+      return;
+    }
+
+    pos_ = 0;
+    tracks_.clear();
+
+    if (shuffle_) {
+      shuffle_->resize(0);
     }
   }
 
   notifyChanged(true);
 }
 
-auto TrackQueue::clear() -> void {
-  const std::unique_lock<std::shared_mutex> lock(mutex_);
-  pos_ = 0;
-  tracks_.clear();
+auto TrackQueue::random(bool en) -> void {
+  {
+    const std::unique_lock<std::shared_mutex> lock(mutex_);
+    // Don't check for en == true already; this has the side effect that
+    // repeated calls with en == true will re-shuffle.
+    if (en) {
+      shuffle_.emplace(tracks_.size());
+      shuffle_->repeat(repeat_);
+    } else {
+      shuffle_.reset();
+    }
+  }
 
-  notifyChanged(true);
+  // Current track doesn't get randomised until next().
+  notifyChanged(false);
+}
+
+auto TrackQueue::random() const -> bool {
+  const std::shared_lock<std::shared_mutex> lock(mutex_);
+  return shuffle_.has_value();
+}
+
+auto TrackQueue::repeat(bool en) -> void {
+  {
+    const std::unique_lock<std::shared_mutex> lock(mutex_);
+    repeat_ = en;
+    if (shuffle_) {
+      shuffle_->repeat(en);
+    }
+  }
+
+  notifyChanged(false);
+}
+
+auto TrackQueue::repeat() const -> bool {
+  const std::shared_lock<std::shared_mutex> lock(mutex_);
+  return repeat_;
 }
 
 }  // namespace audio
diff --git a/src/lua/stubs/queue.lua b/src/lua/stubs/queue.lua
new file mode 100644
index 00000000..000c35d3
--- /dev/null
+++ b/src/lua/stubs/queue.lua
@@ -0,0 +1,22 @@
+--- Properties and functions for inspecting and manipulating the track playback queue
+-- @module queue
+
+local queue = {}
+
+--- queue.position returns the index in the queue of the currently playing track. This may be zero if the queue is empty.
+-- @treturn types.Property a positive integer property, which is a 1-based index
+function queue.position() end
+
+--- queue.size returns the total number of tracks in the queue, including tracks which have already been played.
+-- @treturn types.Property a positive integer property
+function queue.size() end
+
+--- queue.replay determines whether or not the queue will be restarted after the final track is played.
+-- @treturn types.Property a writeable boolean property
+function queue.replay() end
+
+--- queue.random determines whether, when progressing to the next track in the queue, the next track will be chosen randomly. The random selection algorithm used is a Miller Shuffle, which guarantees that no repeat selections will be made until every item in the queue has been played.
+-- @treturn types.Property a writeable boolean property
+function queue.random() end
+
+return queue
\ No newline at end of file
diff --git a/src/ui/include/ui_fsm.hpp b/src/ui/include/ui_fsm.hpp
index 7ad6be93..ba3f5e3f 100644
--- a/src/ui/include/ui_fsm.hpp
+++ b/src/ui/include/ui_fsm.hpp
@@ -137,6 +137,8 @@ class Lua : public UiState {
   auto PushLuaScreen(lua_State*) -> int;
   auto PopLuaScreen(lua_State*) -> int;
   auto SetPlaying(const lua::LuaValue&) -> bool;
+  auto SetRandom(const lua::LuaValue&) -> bool;
+  auto SetRepeat(const lua::LuaValue&) -> bool;
 
   std::shared_ptr<lua::Property> battery_pct_;
   std::shared_ptr<lua::Property> battery_mv_;
@@ -150,6 +152,8 @@ class Lua : public UiState {
 
   std::shared_ptr<lua::Property> queue_position_;
   std::shared_ptr<lua::Property> queue_size_;
+  std::shared_ptr<lua::Property> queue_repeat_;
+  std::shared_ptr<lua::Property> queue_random_;
 };
 
 class Browse : public UiState {
diff --git a/src/ui/ui_fsm.cpp b/src/ui/ui_fsm.cpp
index 783494dd..f4b56a27 100644
--- a/src/ui/ui_fsm.cpp
+++ b/src/ui/ui_fsm.cpp
@@ -176,6 +176,8 @@ void Lua::entry() {
 
     queue_position_ = std::make_shared<lua::Property>(0);
     queue_size_ = std::make_shared<lua::Property>(0);
+    queue_repeat_ = std::make_shared<lua::Property>(false);
+    queue_random_ = std::make_shared<lua::Property>(false);
 
     playback_playing_ = std::make_shared<lua::Property>(
         false, [&](const lua::LuaValue& val) { return SetPlaying(val); });
@@ -203,6 +205,8 @@ void Lua::entry() {
     sLua->bridge().AddPropertyModule("queue", {
                                                   {"position", queue_position_},
                                                   {"size", queue_size_},
+                                                  {"replay", queue_repeat_},
+                                                  {"random", queue_random_},
                                               });
     sLua->bridge().AddPropertyModule(
         "backstack",
@@ -242,7 +246,7 @@ auto Lua::PushLuaScreen(lua_State* s) -> int {
   return 0;
 }
 
-auto Lua::PopLuaScreen(lua_State *s) -> int {
+auto Lua::PopLuaScreen(lua_State* s) -> int {
   PopScreen();
   luavgl_set_root(s, sCurrentScreen->content());
   lv_group_set_default(sCurrentScreen->group());
@@ -261,6 +265,24 @@ auto Lua::SetPlaying(const lua::LuaValue& val) -> bool {
   return true;
 }
 
+auto Lua::SetRandom(const lua::LuaValue& val) -> bool {
+  if (!std::holds_alternative<bool>(val)) {
+    return false;
+  }
+  bool b = std::get<bool>(val);
+  sServices->track_queue().random(b);
+  return true;
+}
+
+auto Lua::SetRepeat(const lua::LuaValue& val) -> bool {
+  if (!std::holds_alternative<bool>(val)) {
+    return false;
+  }
+  bool b = std::get<bool>(val);
+  sServices->track_queue().repeat(b);
+  return true;
+}
+
 void Lua::exit() {
   lv_group_set_default(NULL);
 }
@@ -288,6 +310,8 @@ void Lua::react(const audio::QueueUpdate&) {
     current_pos++;
   }
   queue_position_->Update(current_pos);
+  queue_random_->Update(queue.random());
+  queue_repeat_->Update(queue.repeat());
 }
 
 void Lua::react(const audio::PlaybackStarted& ev) {