Store horizon along with visited nodes

When a node is reached at horizon 0 it is marked as visited.
After reaching the node again, with horizon >0 its relationships
should be visited.

This commit stores the horizon in CypherContext and uses it to evaluate
if the traversal should continue.

Fixes #407.

CHANGES.txt

@@ -1,6 +1,7 @@
 3.0.2
 o Entity count returns incorrect result on abstract non-annotated type. #435
 o Fix classpath scanning issue with Play framework. #429
+o Store horizon along with visited nodes to traverse to correct depth. #407
 
 3.0.1
 o Add filter function for in-collection query. #423

core/src/main/java/org/neo4j/ogm/context/EntityGraphMapper.java

@@ -208,14 +208,13 @@ private NodeBuilder mapEntity(Object entity, int horizon, Compiler compiler) {
             return null;
         }
 
-        if (context.visited(entity)) {
+        if (context.visited(entity, horizon)) {
             LOGGER.debug("already visited: {}", entity);
             return context.visitedNode(entity);
         }
 
-        NodeBuilder nodeBuilder = getNodeBuilder(compiler, entity);
+        NodeBuilder nodeBuilder = getNodeBuilder(compiler, entity, horizon);
         if (nodeBuilder != null) {
-            updateNode(entity, context, nodeBuilder);
             if (horizon != 0) {
                 mapEntityReferences(entity, nodeBuilder, horizon - 1, compiler);
             } else {
@@ -258,9 +257,16 @@ private void updateNode(Object entity, CompileContext context, NodeBuilder nodeB
      *
      * @param compiler the {@link org.neo4j.ogm.cypher.compiler.Compiler}
      * @param entity the object to save
+     * @param horizon current horizon
      * @return a {@link NodeBuilder} object for either a new node or an existing one
      */
-    private NodeBuilder getNodeBuilder(Compiler compiler, Object entity) {
+    private NodeBuilder getNodeBuilder(Compiler compiler, Object entity, int horizon) {
+        CompileContext context = compiler.context();
+
+        NodeBuilder nodeBuilder = context.visitedNode(entity);
+        if (nodeBuilder != null) {
+            return nodeBuilder;
+        }
 
         ClassInfo classInfo = metaData.classInfo(entity);
 
@@ -269,23 +275,23 @@ private NodeBuilder getNodeBuilder(Compiler compiler, Object entity) {
             return null;
         }
 
-        CompileContext context = compiler.context();
         Long id = mappingContext.nativeId(entity);
         Collection<String> labels = EntityUtils.labels(entity, metaData);
 
-        NodeBuilder nodeBuilder;
         final String primaryIndex = classInfo.primaryIndexField() != null ? classInfo.primaryIndexField().property() : null;
         if (id < 0) {
-//            Long entityIdRef = EntityUtils.identity(entity, metaData);
             nodeBuilder = compiler.newNode(id).addLabels(labels).setPrimaryIndex(primaryIndex);
             context.registerNewObject(id, entity);
         } else {
             nodeBuilder = compiler.existingNode(Long.valueOf(id.toString()));
             nodeBuilder.addLabels(labels).setPrimaryIndex(primaryIndex);
             removePreviousLabelsIfRequired(entity, classInfo, nodeBuilder);
         }
-        context.visit(entity, nodeBuilder);
+        context.visit(entity, nodeBuilder, horizon);
         LOGGER.debug("visiting: {}", entity);
+
+        updateNode(entity, context, nodeBuilder);
+
         return nodeBuilder;
     }
 
@@ -585,15 +591,15 @@ private void mapRelationshipEntity(Object relationshipEntity, Object parent, Rel
         NodeBuilder tgtNodeBuilder = context.visitedNode(targetEntity);
 
         if (parent == targetEntity) {  // we approached this RE from its END-NODE during object mapping.
-            if (!context.visited(startEntity)) { // skip if we already visited the START_NODE
+            if (!context.visited(startEntity, horizon)) { // skip if we already visited the START_NODE
                 relNodes.source = targetEntity; // set up the nodes to link
                 relNodes.target = startEntity;
                 mapRelatedEntity(cypherCompiler, nodeBuilder, relationshipBuilder, horizon, relNodes);
             } else {
                 updateRelationship(context, tgtNodeBuilder, srcNodeBuilder, relationshipBuilder, relNodes);
             }
         } else { // we approached this RE from its START_NODE during object mapping.
-            if (!context.visited(targetEntity)) {  // skip if we already visited the END_NODE
+            if (!context.visited(targetEntity, horizon)) {  // skip if we already visited the END_NODE
                 relNodes.source = startEntity;  // set up the nodes to link
                 relNodes.target = targetEntity;
                 mapRelatedEntity(cypherCompiler, nodeBuilder, relationshipBuilder, horizon, relNodes);

core/src/main/java/org/neo4j/ogm/cypher/compiler/CompileContext.java

@@ -30,7 +30,7 @@ public interface CompileContext {
 
     boolean removeRegisteredRelationship(Mappable mappable);
 
-    boolean visited(Object entity);
+    boolean visited(Object entity, int horizon);
 
     NodeBuilder visitedNode(Object entity);
 
@@ -44,7 +44,13 @@ public interface CompileContext {
 
     Collection<Object> registry();
 
-    void visit(Object entity, NodeBuilder nodeBuilder);
+    /**
+     * Stores nodeBuilder for given entity with horizon
+     *
+     * if the nodeBuilder for the entity is already present it will be overwritten (or the horizon will change)
+     * the caller should ensure it doesn't happen
+     */
+    void visit(Object entity, NodeBuilder nodeBuilder, int horizon);
 
     boolean visitedRelationshipEntity(Long relationshipIdentity);
 

core/src/main/java/org/neo4j/ogm/cypher/compiler/CypherContext.java

@@ -37,7 +37,7 @@
  */
 public class CypherContext implements CompileContext {
 
-    private final Map<Object, NodeBuilder> visitedObjects = new IdentityHashMap<>();
+    private final Map<Object, NodeBuilderHorizonPair> visitedObjects = new IdentityHashMap<>();
     private final Set<Long> visitedRelationshipEntities = new HashSet<>();
 
     private final Map<Long, Object> createdObjectsWithId = new HashMap<>();
@@ -54,13 +54,14 @@ public CypherContext(Compiler compiler) {
         this.compiler = compiler;
     }
 
-    public boolean visited(Object entity) {
-        return this.visitedObjects.containsKey(entity);
+    public boolean visited(Object entity, int horizon) {
+        NodeBuilderHorizonPair pair = visitedObjects.get(entity);
+        return pair != null && pair.getHorizon() > horizon;
     }
 
     @Override
-    public void visit(Object entity, NodeBuilder nodeBuilder) {
-        this.visitedObjects.put(entity, nodeBuilder);
+    public void visit(Object entity, NodeBuilder nodeBuilder, int horizon) {
+        this.visitedObjects.put(entity, new NodeBuilderHorizonPair(nodeBuilder, horizon));
     }
 
     public void registerRelationship(Mappable mappedRelationship) {
@@ -73,7 +74,8 @@ public boolean removeRegisteredRelationship(Mappable mappedRelationship) {
 
     @Override
     public NodeBuilder visitedNode(Object entity) {
-        return this.visitedObjects.get(entity);
+        NodeBuilderHorizonPair pair = this.visitedObjects.get(entity);
+        return pair != null ? pair.getNodeBuilder() : null;
     }
 
     @Override
@@ -265,4 +267,23 @@ private boolean isMappableAlreadyDeleted(Mappable mappedRelationship) {
         }
         return false;
     }
+
+    private static class NodeBuilderHorizonPair {
+
+        private final NodeBuilder nodeBuilder;
+        private final int horizon;
+
+        private NodeBuilderHorizonPair(NodeBuilder nodeBuilder, int horizon) {
+            this.nodeBuilder = nodeBuilder;
+            this.horizon = horizon;
+        }
+
+        public NodeBuilder getNodeBuilder() {
+            return nodeBuilder;
+        }
+
+        public int getHorizon() {
+            return horizon;
+        }
+    }
 }