lib/universe_compiler/universe/entities.rb in universe_compiler-0.5.1 vs lib/universe_compiler/universe/entities.rb in universe_compiler-0.5.2

@@ -52,10 +52,13 @@
               end
             end
           end
         end
         # Then delete the entity
+        if entity.respond_to? :'universe='
+          entity.universe = nil
+        end
         entities.delete entity
         reindex_all entities
         unless impacted_entities.empty?
           UniverseCompiler.logger.warn "Impacts of '#{entity.as_path}' deletion:"
           impacted_entities.each do |k, v|
@@ -94,10 +97,16 @@
             end
           end
         end
         # Then replace the entity
         entities[entities.index(entity)] = by
+        if entity.respond_to? :'universe='
+          entity.universe = nil
+        end
+        if by.respond_to? :'universe='
+          by.universe = self
+        end
         reindex_all entities
         unless impacted_entities.empty?
           UniverseCompiler.logger.warn "Impacts of '#{entity.as_path}' deletion:"
           impacted_entities.each do |k, v|
             UniverseCompiler.logger.warn " - '#{k.as_path}' fields : #{v.join ', '}"
@@ -105,9 +114,10 @@
         end
         impacted_entities
       end
 
       def clear
+        entities.each { |e| e.universe = nil }
         entities.clear
         clear_indices
       end
 
       private