/* jojodiff.pk - Patch format of JojoDiff (a diff utility for binary files).  */

/* Copyright (C) 2023, 2024 The poke authors.  */

/* This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

/* Reference:
     - https://jojodiff.sourceforge.net/
     - https://github.com/vibhorkalley/jojodiff
     - https://github.com/janjongboom/janpatch (see issue 16)
 */

load ios;

var JOJO_ESC = 0xa7UB,  /* Escape.  */
    JOJO_MOD = 0xa6UB,  /* Modify.  */
    JOJO_INS = 0xa5UB,  /* Insert.  */
    JOJO_DEL = 0xa4UB,  /* Delete.  */
    JOJO_EQL = 0xa3UB,  /* Equal.   */
    JOJO_BKT = 0xa2UB;  /* Backtrack.  */

type Jojo_Offset = offset<uint<64>,B>;

type Jojo_Length1 =
  struct
  {
    uint<8> x : x < 252UB;

    computed Jojo_Offset value;

    method get_value = Jojo_Offset:
      { return (1UL + x)#B; }

    method _print = void:
      { printf ("#<%v>", get_value); }
  };

type Jojo_Length2 =
  struct
  {
    uint<8> prefix == 252UB;
    uint<8> x;

    computed Jojo_Offset value;

    method get_value = Jojo_Offset: { return (253UL + x)#B; }

    method _print = void:
      { printf ("#<%v>", get_value); }
  };

type Jojo_Length3 =
  struct
  {
    uint<8> prefix == 253UB;
    big uint<16> x;

    computed Jojo_Offset value;

    method get_value = Jojo_Offset: { return x#B; }

    method _print = void:
      { printf ("#<%v>", get_value); }
  };

type Jojo_Length5 =
  struct
  {
    uint<8> prefix == 254UB;
    big uint<32> x;

    computed Jojo_Offset value;

    method get_value = Jojo_Offset: { return x#B; }

    method _print = void:
      { printf ("#<%v>", get_value); }
  };

type Jojo_Length9 =
  struct
  {
    uint<8> prefix == 255UB;
    big uint<64> x;

    computed Jojo_Offset value;

    method get_value = Jojo_Offset: { return x#B; }

    method _print = void:
      { printf ("#<%v>", get_value); }
  };

type Jojo_Length =
  union
  {
    Jojo_Length1 len1;
    Jojo_Length2 len2;
    Jojo_Length3 len3;
    Jojo_Length5 len5;
    Jojo_Length9 len9;

    method get_length = Jojo_Offset:
      {
        if (!(len1 ?! E_elem))
          return len1.value;
        else if (!(len2 ?! E_elem))
          return len2.value;
        else if (!(len3 ?! E_elem))
          return len3.value;
        else if (!(len5 ?! E_elem))
          return len5.value;
        else if (!(len9 ?! E_elem))
          return len9.value;
        else
          assert (0, "unreachable reached!");
      }
  };

/* Encode the length in Jojo format.
 *
 *              0     <  x <= 252           1 byte:   0-251
 *            252     <  x <= 508           2 bytes:  252, x-253
 *            508     <  x <  0x10000       3 bytes:  253, xx
 *        0x10000     <= x <  0x100000000   5 bytes:  254, xxxx
 *        0x100000000 <= x                  9 bytes:  255, xxxxxxxx
 */

fun jojo_length_encode = (Jojo_Offset len) Jojo_Length:
{
  assert (len != 0#B);

  return len <= 252#B ? Jojo_Length {len1 = Jojo_Length1 {x = len/#B - 1}}
       : len <= 508#B ? Jojo_Length {len2 = Jojo_Length2 {x = len/#B - 253}}
       : len < 0x10000#B ? Jojo_Length {len3 = Jojo_Length3 {x = len/#B}}
       : len < 0x100000000#B ? Jojo_Length {len5 = Jojo_Length5 {x = len/#B}}
       : Jojo_Length {len9 = Jojo_Length9 {x = len/#B}};
}

type Jojo_Datum =
  union
  {
    uint<8>[2] escaped_escape : escaped_escape == [JOJO_ESC, JOJO_ESC];
    uint<8>[2] pair : pair[0] == JOJO_ESC
        && !(pair[1] in [JOJO_MOD, JOJO_INS, JOJO_DEL, JOJO_EQL, JOJO_BKT]);
    uint<8> value : value != JOJO_ESC;
  };

/* Jojo-encoded bytes.  */

type Jojo_Bytes =
  struct
  {
    Jojo_Datum[] data : data'length > 0;

    method get_count = Jojo_Offset:
      {
        var len = 0UL#B;

        for (d in data)
          len += !(d.pair ?! E_elem) ? 2#B : 1#B;
        return len;
      }
  };

type Jojo_Bytes_Callback = (uint<64>/*index*/,uint<8>/*byte*/)void;

/* Traverse Jojo_Bytes's data (Jojo_Datum) and for each byte call provided
   DO callback. */

fun jojo_bytes_map = (Jojo_Bytes jbytes, Jojo_Bytes_Callback do) void:
{
  var i = 0UL;

  for (d in jbytes.data)
    if (!(d.value ?! E_elem))
      do (/*index*/ i++, /*byte*/ d.value);
    else if (!(d.pair ?! E_elem))
      {
        do (/*index*/ i++, /*byte*/ d.pair[0]);
        do (/*index*/ i++, /*byte*/ d.pair[1]);
      }
    else
      do (/*index*/ i++, /*byte*/ JOJO_ESC);
}

/* Decode Jojo-encoded bytes into normal byte array.  */

fun jojo_bytes = (Jojo_Bytes jbytes) uint<8>[]:
{
  var bytes = uint<8>[jbytes.get_count] ();

  jojo_bytes_map
    :jbytes jbytes
    :do lambda (uint<64> i, uint<8> byte) void:
          {
            bytes[i] = byte;
          };
  return bytes;
}

/* Encode provided BYTES in Jojo format.  */

fun jojo_bytes_encode = (uint<8>[] bytes) Jojo_Bytes:
{
  var ds = Jojo_Datum[] ();

  for (b in bytes)
    {
      if (b == JOJO_ESC)
        apush (ds, Jojo_Datum {escaped_escape=[JOJO_ESC, JOJO_ESC]});
      else
        apush (ds, Jojo_Datum {value=b});
    }
  return Jojo_Bytes {data=ds};
}

/* MOD patch instruction

   Add the following bytes to the new file.  Advance both cursors.  */

type Jojo_MOD =
  struct
  {
    uint<8> prefix == JOJO_ESC;
    uint<8> opcode == JOJO_MOD;
    Jojo_Bytes jbytes;

    computed Jojo_Offset count;

    method get_count = Jojo_Offset:
      { return jbytes.get_count; }

    method _print = void:
      { printf ("#<MOD:%v>", get_count); }
  };

/* INS patch instruction

   Add the following bytes to the new file.  Advance cursor in new file.  */

type Jojo_INS =
  struct
  {
    uint<8> prefix == JOJO_ESC;
    uint<8> opcode == JOJO_INS;
    Jojo_Bytes jbytes;

    computed Jojo_Offset count;

    method get_count = Jojo_Offset:
      { return jbytes.get_count; }  /* Don't use `size' attribute.  */

    method _print = void:
      { printf ("#<INS:%v>", get_count); }
  };

/* DEL patch instruction

   Advance cursor in original file.  */

type Jojo_DEL =
  struct
  {
    uint<8> prefix == JOJO_ESC;
    uint<8> opcode == JOJO_DEL;
    Jojo_Length jojo_length;

    computed Jojo_Offset count;

    method get_count = Jojo_Offset:
      { return jojo_length.get_length; }

    method _print = void:
      { printf ("#<DEL:%v>", get_count); }
  };

/* BKT patch instruction

   Backtrack in the original file.  */

type Jojo_BKT =
  struct
  {
    uint<8> prefix == JOJO_ESC;
    uint<8> opcode == JOJO_BKT;
    Jojo_Length jojo_length;

    computed Jojo_Offset count;

    method get_count = Jojo_Offset:
      { return jojo_length.get_length; }

    method _print = void:
      { printf ("#<BKT:%v>", get_count); }
  };

/* EQL patch instruction

   Copy from original file to new file.  Advance both cursors.  */

type Jojo_EQL =
  struct
  {
    uint<8> prefix == JOJO_ESC;
    uint<8> opcode == JOJO_EQL;
    Jojo_Length jojo_length;

    computed Jojo_Offset count;

    method get_count = Jojo_Offset:
      { return jojo_length.get_length; }

    method _print = void:
      { printf ("#<EQL:%v>", get_count); }
  };

/* In commit 75e4450b2ff661b9d3d0222ca84865e4131f714c (jojodiff.sf.net),
   the author changed the format of output in a backward-incompatible way.
   With that commit, ESC-MOD sequence is the default operation when a new
   operation sequence is needed:
     - at the start of a file
     - after an EQL, DEL or BKT operation

   Use Jojo_{Hunk,Patch} for new data, and Jojo_{Hunk,Patch}_Orig for the
   original data.  */

type Jojo_Hunk =
  union
  {
    Jojo_MOD mod;
    Jojo_INS ins;
    Jojo_DEL del;
    Jojo_EQL eql;
    Jojo_BKT bkt;
    Jojo_Bytes jbytes;
  };

type Jojo_Hunk_Orig =
  union
  {
    Jojo_MOD mod;
    Jojo_INS ins;
    Jojo_DEL del;
    Jojo_EQL eql;
    Jojo_BKT bkt;
  };

type Jojo_Patch = Jojo_Hunk[];
type Jojo_Patch_Orig = Jojo_Hunk_Orig[];

/* Convert original Jojo patch to Jojo_Patch type.

   Useful as a pre-processing step before `jojo_patch_apply'.  */

fun jojo_patch_from_orig = (Jojo_Patch_Orig orig) Jojo_Patch:
{
  var hs = Jojo_Hunk[] ();

  for (h in orig)
    {
      if (!(h.mod ?! E_elem))
        apush (hs, Jojo_Hunk {mod=h.mod});
      else if (!(h.ins ?! E_elem))
        apush (hs, Jojo_Hunk {ins=h.ins});
      else if (!(h.del ?! E_elem))
        apush (hs, Jojo_Hunk {del=h.del});
      else if (!(h.eql ?! E_elem))
        apush (hs, Jojo_Hunk {eql=h.eql});
      else if (!(h.bkt ?! E_elem))
        apush (hs, Jojo_Hunk {bkt=h.bkt});
      else
        assert (0, "unreachable reached!");
    }
    return hs;
}

/* Apply provided Jojo PATCH to construct the new content in NEW_IOS from
   content of ORIG_IOS.  */

fun jojo_patch_apply = (Jojo_Patch patch,
                        int<32> orig_ios,
                        int<32> new_ios,
                        Jojo_Offset orig_off = 0#B,
                        Jojo_Offset new_off = 0#B,
                        int<32> verbosity = 0) Jojo_Offset:
{
  for (hunk in patch)
    {
      if (verbosity > 0)
        {
          var oacutoff = vm_oacutoff;

          if (verbosity == 1) vm_set_oacutoff (5);
          printf ("Hunk:%v\n", hunk);
          if (verbosity == 1) vm_set_oacutoff (oacutoff);
        }

      if (!(hunk.mod ?! E_elem))
        {
          var jbytes = hunk.mod.jbytes;

          if (verbosity > 1)
            {
              printf ("MOD :to %v :from %v", new_off, orig_off);
              if (verbosity > 2)
                {
                  var bytes = jojo_bytes (jbytes);

                  printf (" :nbytes %v :bytes %v", bytes'size, bytes);
                }
              print ("\n");
            }

          /* Add the `bytes' to the new file.
             Advance both cursors.  */
          jojo_bytes_map
            :jbytes jbytes
            :do lambda (uint<64> index, uint<8> byte) void:
                {
                  uint<8> @ new_ios : new_off = byte;
                  new_off++;
                  orig_off++;
                };
        }
      else if (!(hunk.ins ?! E_elem))
        {
          var jbytes = hunk.ins.jbytes;

          if (verbosity > 1)
            {
              printf ("INS :to %v", new_off);
              if (verbosity > 2)
                {
                  var bytes = jojo_bytes (jbytes);

                  printf (" :nbytes %v :bytes %v", bytes'size, bytes);
                }
              print ("\n");
            }

          /* Add the following bytes to the new file.
             Advance cursor in new file.  */
          jojo_bytes_map
            :jbytes jbytes
            :do lambda (uint<64> index, uint<8> byte) void:
                {
                  uint<8> @ new_ios : new_off = byte;
                  new_off++;
                };
        }
      else if (!(hunk.del ?! E_elem))
        {
          var len = hunk.del.jojo_length.get_length;

          if (verbosity > 1)
            {
              printf ("DEL :from %v :size %v\n", orig_off, len);
              if (verbosity > 2)
                ios_dump_bytes :ios orig_ios
                               :from orig_off
                               :size len
                               :ascii_p 1;
            }

          orig_off += len;
        }
      else if (!(hunk.eql ?! E_elem))
        {
          var len = hunk.eql.jojo_length.get_length;

          if (verbosity > 1)
            {
              printf ("EQL :from %v :size %v\n", orig_off, len);
              ios_dump_bytes :ios orig_ios
                             :from orig_off
                             :size len
                             :ascii_p 1;
            }

          ios_copy_bytes :from_ios orig_ios :to_ios new_ios
                         :from orig_off :to new_off :size len;
          orig_off += len;
          new_off += len;
        }
      else if (!(hunk.bkt ?! E_elem))
        {
          var len = hunk.bkt.jojo_length.get_length;

          if (verbosity > 1)
            printf ("BKT :from %v :len %v\n", orig_off, len);

          assert (orig_off >= len);
          orig_off -= len;
        }
      else if (!(hunk.jbytes ?! E_elem))
        {
          /* Like JOJO_Mod.  */
          var jbytes = hunk.jbytes;

          if (verbosity > 1)
            {
              printf ("MOD(implicit) :to %v :from %v", new_off, orig_off);
              if (verbosity > 2)
                {
                  var bytes = jojo_bytes (jbytes);

                  printf (" :nbytes %v :bytes %v", bytes'size, bytes);
                }
              print ("\n");
            }

          /* Add the `bytes' to the new file.
             Advance both cursors.  */
          jojo_bytes_map
            :jbytes jbytes
            :do lambda (uint<64> index, uint<8> byte) void:
                {
                  uint<8> @ new_ios : new_off = byte;
                  new_off++;
                  orig_off++;
                };
        }
      else
        assert (0, "unreachable reached!");
    }
  return new_off;
}
