3 files changed, 9 insertions, 6 deletions
diff --git a/src/script.cpp b/src/script.cpp
index e3aca3e..f52391f 100644
--- a/src/script.cpp
+++ b/src/script.cpp
@@ -90,7 +90,7 @@ void ScriptSystem::scriptExport(void)
     std::function<sol::table(sol::table)> entitySpawn = 
         [this](sol::table t){ return spawn(t);};
 
-    std::function<unsigned int(sol::object)> worldRegister =
+    std::function<World* (sol::object)> worldRegister =
         [this](sol::object t){ return worldSystem.addWorld(t); };
 
     lua.new_usertype<Position>("Position",
diff --git a/src/world.cpp b/src/world.cpp
index 8b71316..63769a9 100644
--- a/src/world.cpp
+++ b/src/world.cpp
@@ -38,10 +38,12 @@ World::World(sol::object ref)
 *  WORLD SYSTEM  *
 ******************/
 
-unsigned int WorldSystem::addWorld(sol::object t)
+World* WorldSystem::addWorld(sol::object t)
 {
     worlds.push_back(World(t));
-    return worlds.size()-1;
+    if (currentWorld == nullptr)
+        currentWorld = &(worlds.back());
+    return &(worlds.back());
 }
 
 void WorldSystem::configure([[maybe_unused]]entityx::EntityManager& entities,
diff --git a/src/world.hpp b/src/world.hpp
index 7c66e6b..4ceca94 100644
--- a/src/world.hpp
+++ b/src/world.hpp
@@ -52,13 +52,14 @@ class WorldSystem : public entityx::System<WorldSystem>
 {
 private:
     std::vector<World> worlds;
-    //World& currentWorld;
+    World* currentWorld;
 public:
-    WorldSystem(void) {}
+    WorldSystem(void):
+        currentWorld(nullptr) {}
 
     ~WorldSystem(void) {}
 
-    unsigned int addWorld(sol::object);
+    World* addWorld(sol::object);
 
     /**
      * Prepares the system for running.