From e3ab67aed48b5c05ae0ac3f107df70262a7c9552 Mon Sep 17 00:00:00 2001
From: Michael Munday <munday@ca.ibm.com>
Date: Tue, 6 Dec 2016 11:56:37 -0500
Subject: [PATCH] Support forward local branches.

---
 dynasm/Examples/test_z_inst.c | 18 ++++++++++++++++++
 dynasm/dasm_s390x.h           | 18 ++++++++++++++----
 dynasm/dasm_s390x.lua         |  9 ++++++++-
 3 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/dynasm/Examples/test_z_inst.c b/dynasm/Examples/test_z_inst.c
index 547a1c46..3938a741 100644
--- a/dynasm/Examples/test_z_inst.c
+++ b/dynasm/Examples/test_z_inst.c
@@ -76,6 +76,23 @@ static void lab(dasm_State *state)
   | br r14
 }
 
+static void labg(dasm_State *state)
+{
+  dasm_State **Dst = &state;
+
+  // r1 = 0; do { r2 += r2; r1 += 1; } while(r1 < r3);
+  | la r1, 0(r0)
+  |1:
+  | agr r2, r2
+  | la r1, 1(r1)
+  | cgr r1, r3
+  | jgl <1
+  | jgnl >1
+  | stg r0, 0(r0)
+  |1:
+  | br r14
+}
+
 static void add_imm16(dasm_State *state)
 {
   dasm_State **Dst = &state;
@@ -108,6 +125,7 @@ test_table test[] = {
   { 5, 7,        rx,  12298,    "rx"},
   { 5, 7,       rxy,     10,   "rxy"},
   { 2, 4,       lab,     32,   "lab"},
+  { 2, 4,      labg,     32,  "labg"},
   { 2, 0, add_imm16,     17, "imm16"},
   { 2, 0, add_imm32,     16, "imm32"}
 };
diff --git a/dynasm/dasm_s390x.h b/dynasm/dasm_s390x.h
index 8fc4cd63..18a7338e 100644
--- a/dynasm/dasm_s390x.h
+++ b/dynasm/dasm_s390x.h
@@ -210,7 +210,7 @@ void dasm_put(Dst_DECL, int start, ...)
       goto stop;
     case DASM_ESC:
       p++;
-      ofs += 4;
+      ofs += 2;
       break;
     case DASM_REL_EXT:
       break;
@@ -244,6 +244,10 @@ void dasm_put(Dst_DECL, int start, ...)
         b[pos] = n;             /* Else link to rel chain, anchored at label. */
         *pl = pos;
       }
+      ofs += 2;
+      if (p[-3] >> 12 == 0xc) { /* RIL instruction needs 32-bit immediate. */
+        ofs += 2;
+      }
       pos++;
       break;
     case DASM_LABEL_LG:
@@ -262,7 +266,6 @@ void dasm_put(Dst_DECL, int start, ...)
       }
       *pl = -pos;               /* Label exists now. */
       b[pos++] = ofs;           /* Store pass1 offset estimate. */
-      ofs += 2;
       break;
     case DASM_IMM16:
       CK(((short)n) == n, RANGE_I);     /* TODO: unsigned immediates? */
@@ -424,8 +427,15 @@ int dasm_encode(Dst_DECL, void *buffer)
           CK(n >= 0, UNDEF_PC);
           n = *DASM_POS2PTR(D, n) - (int)((char *)cp - base);
         patchrel:
-          *cp++ = n/2; /* TODO: only 16-bit relative jump currently works. */
-          p++; /* skip argument */
+          /* Offsets are halfword aligned (so need to be halved). */
+          n += 2;               /* Offset is relative to start of instruction. */
+          if (cp[-1] >> 12 == 0xc) {
+            *cp++ = n >> 17;
+          } else {
+            CK(-(1 << 16) <= n && n < (1 << 16) && n & 1 == 0, RANGE_LG);
+          }
+          *cp++ = n >> 1;
+          p++;                  /* skip argument */
           break;
         case DASM_LABEL_LG:
           ins = *p++;
diff --git a/dynasm/dasm_s390x.lua b/dynasm/dasm_s390x.lua
index 9a77d4e6..847a02e7 100644
--- a/dynasm/dasm_s390x.lua
+++ b/dynasm/dasm_s390x.lua
@@ -75,6 +75,13 @@ local function dumpactions(out)
   out:write("\n")
 end
 
+local function havearg(a)
+  return a == "ESC" or
+         a == "SECTION" or
+         a == "REL_LG" or
+         a == "LABEL_LG"
+end
+
 -- Write action list buffer as a huge static C array.
 local function writeactions(out, name)
   local nn = #actlist
@@ -87,7 +94,7 @@ local function writeactions(out, name)
     local name = action_names[actlist[i]+1]
     if not esc and name then
       assert(out:write(" /* ", name, " */"))
-      esc = name == "ESC" or name == "SECTION"
+      esc = havearg(name)
     else
       esc = false
     end