diff --git a/includes/control_menu.hpp b/includes/control_menu.hpp
index 1660da7..2aa0c2e 100644
--- a/includes/control_menu.hpp
+++ b/includes/control_menu.hpp
@@ -9,6 +9,7 @@
 #pragma once
 
 #include <context.hpp>
+#include <types.hpp>
 
 namespace gol {
 
@@ -21,6 +22,7 @@ public:
 private:
   ctx context_;
 public: // Keep those public for easy access
+  MenuState menu_state_ = MenuState::NONE;
   int fps_ctrl_ = false;
   int cell_size_ctrl_ = false;
   bool play_ctrl_ = true;
diff --git a/src/control_menu.cpp b/src/control_menu.cpp
index 0f10f0d..dfb3fc5 100644
--- a/src/control_menu.cpp
+++ b/src/control_menu.cpp
@@ -6,7 +6,12 @@
 * ------
 */
 
+#include <imgui.h>
+#include <raylib.h>
+#include <rlImGui.h>
+
 #include <control_menu.hpp>
+#include <world.hpp>
 
 namespace gol {
 
@@ -27,4 +32,59 @@ ControlMenu::ControlMenu(ctx context):context_(context) {
   apply_ctrl_ = false;
 }
 
+void ControlMenu::update() {
+  if (menu_state_ == MenuState::PLAY || menu_state_ == MenuState::EDIT) {
+    if (!play_ctrl_ && !edit_ctrl_) {
+      menu_state_ = MenuState::NONE;
+    }
+  }
+  if (rand_ctrl_) {
+    context_.world->randomize();
+    rand_ctrl_ = false;
+  }
+  if (clear_ctrl_) {
+    context_.world->clear();
+    clear_ctrl_ = false;
+  }
+  if (edit_ctrl_ && play_ctrl_) {
+    if (menu_state_ == MenuState::PLAY) {
+      menu_state_ = MenuState::EDIT;
+      play_ctrl_ = false;
+    } else if (menu_state_ == MenuState::EDIT) {
+      menu_state_ = MenuState::PLAY;
+      edit_ctrl_ = false;
+    }
+  } else if (play_ctrl_) {
+    menu_state_ = MenuState::PLAY;
+  } else if (edit_ctrl_) {
+    menu_state_ = MenuState::EDIT;
+  }
+}
+
+void ControlMenu::display() {
+    bool control_panel_open = true;
+    ImGuiWindowFlags control_panel_flags = ImGuiWindowFlags_AlwaysAutoResize |
+                                           ImGuiWindowFlags_NoResize;
+    ImVec2 menu_pos(3, 3);
+    ImGui::SetNextWindowPos(menu_pos, ImGuiCond_Always);
+    ImGui::Begin("Control Panel", &control_panel_open, control_panel_flags);
+    ImGui::Checkbox("Play", &play_ctrl_);
+    if (ImGui::Button("Step")) {
+      step_ctrl_ = true;
+    }
+    if (ImGui::Button("Step Back")) {
+      step_back_ctrl_ = true;
+    }
+    ImGui::Checkbox("Edit", &edit_ctrl_);
+    ImGui::Checkbox("Clear", &clear_ctrl_);
+    ImGui::Checkbox("Randomize", &rand_ctrl_);
+    if (ImGui::Button((paterns_ctrl_) ? "Hide paterns" : "Show paterns")) {
+      paterns_ctrl_ = !paterns_ctrl_;
+    }
+    if (ImGui::Button((settings_window_) ? "Hide settings" : "Show settings")) {
+      settings_window_ = !settings_window_;
+    }
+    ImGui::Text("Generation: %zu", context_.world->getCycle());
+    ImGui::End();}
+
 } // namespace gol
diff --git a/src/main.cpp b/src/main.cpp
index 45f8a3f..cbbaa43 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -37,7 +37,6 @@ int main(int ac, char **av) {
   gol::ctx context;
   // Load or default config
   const std::string config_file_name = "config.json";
-  MenuState menu_state = MenuState::NONE;
 
   std::fstream config_file(config_file_name, std::ios::in | std::ios::out);
 
@@ -126,45 +125,22 @@ int main(int ac, char **av) {
   rules.setup(&(*context.world));
   // Diplay generations
   while (!WindowShouldClose()) {
-    if (menu_state == MenuState::PLAY || menu_state == MenuState::EDIT) {
-      if (!control_menu.play_ctrl_ && !control_menu.edit_ctrl_) {
-        menu_state = MenuState::NONE;
-      }
-    }
     // Frames shinenigans
     float deltaTime = GetFrameTime();
     sim_speed = control_menu.fps_ctrl_ / 10.0f;
     timePerUpdate = (1.0f / 10.0f) / sim_speed;
 
+    control_menu.update();
+
     auto gesture = GetGestureDetected();
     auto mousePos = GetMousePosition();
-    if (control_menu.rand_ctrl_) {
-      context.world->randomize();
-      control_menu.rand_ctrl_ = false;
-    }
-    if (control_menu.clear_ctrl_) {
-      context.world->clear();
-      control_menu.clear_ctrl_ = false;
-    }
-    if (control_menu.edit_ctrl_ && control_menu.play_ctrl_) {
-      if (menu_state == MenuState::PLAY) {
-        menu_state = MenuState::EDIT;
-        control_menu.play_ctrl_ = false;
-      } else if (menu_state == MenuState::EDIT) {
-        menu_state = MenuState::PLAY;
-        control_menu.edit_ctrl_ = false;
-      }
-    } else if (control_menu.play_ctrl_) {
-      menu_state = MenuState::PLAY;
-    } else if (control_menu.edit_ctrl_) {
-      menu_state = MenuState::EDIT;
-    }
     if (control_menu.edit_ctrl_ && IsMouseButtonPressed(0) &&
         !ImGui::IsWindowHovered(ImGuiHoveredFlags_AnyWindow)) {
-      menu_state = MenuState::EDIT;
+      control_menu.menu_state_ = MenuState::EDIT;
       context.world->setCell(mousePos.x / context.config_json["cell_size"].get<int>(),
                     mousePos.y / context.config_json["cell_size"].get<int>());
     }
+
     // Selection behaviour
     if (!control_menu.edit_ctrl_ && !control_menu.play_ctrl_ && !control_menu.paterns_ctrl_ && !control_menu.settings_window_) {
       if (gesture == GESTURE_TAP && !ImGui::IsWindowHovered(ImGuiHoveredFlags_AnyWindow)) {
@@ -240,7 +216,7 @@ int main(int ac, char **av) {
     if (deltaTimeAccumulator >= timePerUpdate) {
       // Reset accumulator
       deltaTimeAccumulator -= timePerUpdate;
-      if (menu_state == MenuState::PLAY) {
+      if (control_menu.menu_state_ == MenuState::PLAY) {
         context.world->saveCompressed();
         rules.update();
       }
@@ -259,31 +235,7 @@ int main(int ac, char **av) {
     }
     // Start ImGui frame
     rlImGuiBegin();
-
-    bool control_panel_open = true;
-    ImGuiWindowFlags control_panel_flags = ImGuiWindowFlags_AlwaysAutoResize |
-                                           ImGuiWindowFlags_NoResize;
-    ImVec2 menu_pos(3, 3);
-    ImGui::SetNextWindowPos(menu_pos, ImGuiCond_Always);
-    ImGui::Begin("Control Panel", &control_panel_open, control_panel_flags);
-    ImGui::Checkbox("Play", &control_menu.play_ctrl_);
-    if (ImGui::Button("Step")) {
-      control_menu.step_ctrl_ = true;
-    }
-    if (ImGui::Button("Step Back")) {
-      control_menu.step_back_ctrl_ = true;
-    }
-    ImGui::Checkbox("Edit", &control_menu.edit_ctrl_);
-    ImGui::Checkbox("Clear", &control_menu.clear_ctrl_);
-    ImGui::Checkbox("Randomize", &control_menu.rand_ctrl_);
-    if (ImGui::Button((control_menu.paterns_ctrl_) ? "Hide paterns" : "Show paterns")) {
-      control_menu.paterns_ctrl_ = !control_menu.paterns_ctrl_;
-    }
-    if (ImGui::Button((control_menu.settings_window_) ? "Hide settings" : "Show settings")) {
-      control_menu.settings_window_ = !control_menu.settings_window_;
-    }
-    ImGui::Text("Generation: %zu", context.world->getCycle());
-    ImGui::End();
+    control_menu.display();
     if (control_menu.paterns_ctrl_) {
       ImGuiWindowFlags paterns_flags =
           ImGuiWindowFlags_AlwaysAutoResize | ImGuiWindowFlags_NoResize;