Updated focus for relation queries.

Change-Id: Id8ca54ae3b874b45165791b34fb2aa1c9b57a21a
diff --git a/src/main/java/de/ids_mannheim/korap/KrillQuery.java b/src/main/java/de/ids_mannheim/korap/KrillQuery.java
index 45290e8..5a0ccd7 100644
--- a/src/main/java/de/ids_mannheim/korap/KrillQuery.java
+++ b/src/main/java/de/ids_mannheim/korap/KrillQuery.java
@@ -85,6 +85,10 @@
 
     private static final int MAX_CLASS_NUM = 255; // 127;
 
+    // Variables used for relation queries
+    private String direction;
+    private byte[] classNumbers;
+    
     // Private class for koral:boundary objects
     private class Boundary {
         public int min, max;
@@ -129,6 +133,10 @@
      */
     public KrillQuery (String field) {
         this.field = field;
+        this.direction = ">:";
+        this.classNumbers = new byte[2];
+        this.classNumbers[0] = (byte) 1;
+        this.classNumbers[1] = (byte) 2;
     };
 
 
@@ -177,10 +185,6 @@
         return this.fromJson(jsonN);
     };
 
-    public SpanQueryWrapper fromJson(JsonNode json) throws QueryException {
-        return fromJson(json, null);
-    }
-
     /**
      * <p>Deserialize JSON-LD query as a {@link JsonNode} object
      * to a {@link SpanQueryWrapper} object.</p>
@@ -194,8 +198,7 @@
     // TODO: Use the shortcuts implemented in the builder
     //       instead of the wrapper constructors
     // TODO: Rename this span context!
-    public SpanQueryWrapper fromJson(JsonNode json, String direction)
-            throws QueryException {
+    public SpanQueryWrapper fromJson(JsonNode json) throws QueryException {
         int number = 0;
 
         // Only accept @typed objects for the moment
@@ -437,21 +440,17 @@
                     "Number of operands is not acceptable");
         }
 
-        String direction = null;
         SpanQueryWrapper operand1 = fromJson(operands.get(0));
         SpanQueryWrapper operand2 = fromJson(operands.get(1));
 
         if (operand1.isEmpty()) {
             direction = "<:";
         }
-        else {// if (operand2.isEmpty()) {
-            direction = ">:";
-        }
 
-        SpanQueryWrapper relationWrapper = fromJson(relation, direction);
+        SpanQueryWrapper relationWrapper = fromJson(relation);
 
-        return new SpanRelationWrapper(relationWrapper,
-                operand1, operand2);
+        return new SpanRelationWrapper(relationWrapper, operand1, operand2,
+                classNumbers);
 
     }
 
@@ -692,6 +691,9 @@
                 throw new QueryException(709, "Valid class numbers exceeded");
             };
 
+            this.classNumbers[0] = (byte) number;
+            this.classNumbers[1] = (byte) 0;
+
             // Serialize operand
             SpanQueryWrapper sqw = this.fromJson(operands.get(0));
 
diff --git a/src/main/java/de/ids_mannheim/korap/query/wrap/SpanRelationWrapper.java b/src/main/java/de/ids_mannheim/korap/query/wrap/SpanRelationWrapper.java
index 1f1c799..c4dd49b 100644
--- a/src/main/java/de/ids_mannheim/korap/query/wrap/SpanRelationWrapper.java
+++ b/src/main/java/de/ids_mannheim/korap/query/wrap/SpanRelationWrapper.java
@@ -13,9 +13,11 @@
     private SpanQueryWrapper relationQuery;
     private SpanQueryWrapper subQuery1;
     private SpanQueryWrapper subQuery2;
+    private byte[] classNumbers;
 
     public SpanRelationWrapper (SpanQueryWrapper relationWrapper,
-            SpanQueryWrapper operand1, SpanQueryWrapper operand2) {
+            SpanQueryWrapper operand1, SpanQueryWrapper operand2,
+            byte[] classNumbers) {
 
         this.relationQuery = relationWrapper;
         if (relationQuery != null) {
@@ -31,6 +33,7 @@
 
         this.subQuery1 = operand1;
         this.subQuery2 = operand2;
+        this.classNumbers = classNumbers;
     }
 
     @Override
@@ -44,8 +47,11 @@
         if (sq == null) return null;
 
         ArrayList<Byte> classNumbers = new ArrayList<Byte>();
-        classNumbers.add((byte) 1);
-        classNumbers.add((byte) 2);
+        for (byte c : this.classNumbers) {
+            if (c > 0) {
+                classNumbers.add(c);
+            }
+        }
 
         SpanQuery subq1, subq2;
         if (subQuery1.isEmpty) {
diff --git a/src/test/java/de/ids_mannheim/korap/index/TestRelationIndex.java b/src/test/java/de/ids_mannheim/korap/index/TestRelationIndex.java
index 02787b1..7e28e1e 100644
--- a/src/test/java/de/ids_mannheim/korap/index/TestRelationIndex.java
+++ b/src/test/java/de/ids_mannheim/korap/index/TestRelationIndex.java
@@ -116,7 +116,7 @@
     }
 
 
-    private FieldDocument createFieldDoc2 () {
+    public static FieldDocument createFieldDoc2() {
         FieldDocument fd = new FieldDocument();
         fd.addString("ID", "doc-2");
         fd.addTV(
diff --git a/src/test/java/de/ids_mannheim/korap/query/TestSpanRelationQueryJSON.java b/src/test/java/de/ids_mannheim/korap/query/TestSpanRelationQueryJSON.java
index 9179301..7de20c6 100644
--- a/src/test/java/de/ids_mannheim/korap/query/TestSpanRelationQueryJSON.java
+++ b/src/test/java/de/ids_mannheim/korap/query/TestSpanRelationQueryJSON.java
@@ -82,4 +82,25 @@
         SpanQuery sq = sqwi.toQuery();
         assertEquals("focus([1,2]tokens:<:mate/d:HEAD)", sq.toString());
     }
+
+    public void testFocusSource() throws QueryException {
+        //
+        String filepath = getClass().getResource(
+                "/queries/relation/focus-source.json").getFile();
+        SpanQueryWrapper sqwi = getJSONQuery(filepath);
+        SpanQuery sq = sqwi.toQuery();
+        assertEquals(
+                "focus([1]spanSegment(focus(2: spanSegment(tokens:>:mate/d:HEAD, <tokens:c:s />)), <tokens:c:vp />))",
+                sq.toString());
+    }
+
+    public void testFocusTarget() throws QueryException {
+        String filepath = getClass().getResource(
+                "/queries/relation/match-source.json").getFile();
+        SpanQueryWrapper sqwi = getJSONQuery(filepath);
+        SpanQuery sq = sqwi.toQuery();
+        assertEquals(
+                "focus([2]spanSegment(tokens:>:mate/d:HEAD, <tokens:c:s />))",
+                sq.toString());
+    }
 }
diff --git a/src/test/resources/queries/relation/focus-source.json b/src/test/resources/queries/relation/focus-source.json
new file mode 100644
index 0000000..8ac35c7
--- /dev/null
+++ b/src/test/resources/queries/relation/focus-source.json
@@ -0,0 +1,31 @@
+{
+    "query": {
+        "@type": "korap:reference",
+        "operation": "operation:focus",
+        "classRef": [1],
+        "operands": [{
+            "@type": "korap:group",
+            "operation": "operation:relation",
+            "operands": [
+                {
+                    "@type": "korap:span",
+                    "key": "s"
+                },
+                {
+                    "@type": "korap:span",
+                    "key": "np"
+                }
+            ],
+            "relation": {
+                "@type": "korap:relation",
+                "wrap": {
+                    "@type": "korap:term",
+                    "foundry": "mate",
+                    "layer": "d",
+                    "key": "HEAD"
+                }
+            }
+        }]
+    },
+    "meta": {}
+}
diff --git a/src/test/resources/queries/relation/focus-target.json b/src/test/resources/queries/relation/focus-target.json
new file mode 100644
index 0000000..ce7ae66
--- /dev/null
+++ b/src/test/resources/queries/relation/focus-target.json
@@ -0,0 +1,30 @@
+{
+    "query": {
+        "@type": "korap:reference",
+        "operation": "operation:focus",
+        "classRef": [2],
+        "operands": [{
+            "@type": "korap:group",
+            "operation": "operation:relation",
+            "operands": [
+                {
+                    "@type": "korap:span",
+                    "key": "s"
+                },
+                {
+                    "@type": "korap:token"
+                }
+            ],
+            "relation": {
+                "@type": "korap:relation",
+                "wrap": {
+                    "@type": "korap:term",
+                    "foundry": "mate",
+                    "layer": "d",
+                    "key": "HEAD"
+                }
+            }
+        }]
+    },
+    "meta": {}
+}