From cf517effb370a120b3457b6677ef6496350109bb Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Sat, 18 Sep 2021 12:49:38 +0530
Subject: [PATCH] Wrap more of the librsync API

---
 kittens/transfer/rsync.c   | 63 ++++++++++++++++++++++++++++++++++++--
 kittens/transfer/rsync.pyi | 16 ++++++++++
 2 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/kittens/transfer/rsync.c b/kittens/transfer/rsync.c
index 9cfd9bd47..1776d4389 100644
--- a/kittens/transfer/rsync.c
+++ b/kittens/transfer/rsync.c
@@ -9,7 +9,9 @@
 #include <librsync.h>
 
 #define JOB_CAPSULE "rs_job_t"
+#define SIGNATURE_CAPSULE "rs_signature_t"
 #define IO_BUFFER_SIZE (64u * 1024u)
+static PyObject *RsyncError = NULL;
 
 static void
 free_job_capsule(PyObject *capsule) {
@@ -17,6 +19,12 @@ free_job_capsule(PyObject *capsule) {
     if (job) rs_job_free(job);
 }
 
+static void
+free_sig_capsule(PyObject *capsule) {
+    rs_signature_t *sig = PyCapsule_GetPointer(capsule, SIGNATURE_CAPSULE);
+    if (sig) rs_free_sumset(sig);
+}
+
 static PyObject*
 begin_create_signature(PyObject *self UNUSED, PyObject *args) {
     long long file_size = -1;
@@ -44,10 +52,10 @@ begin_create_signature(PyObject *self UNUSED, PyObject *args) {
 
 static PyObject*
 iter_job(PyObject *self UNUSED, PyObject *args) {
-    PyObject *job_capsule;
     Py_ssize_t input_data_size;
     char *input_data;
     int eof = -1;
+    PyObject *job_capsule;
     if (!PyArg_ParseTuple(args, "O!y#|p", &PyCapsule_Type, &job_capsule, &input_data, &input_data_size, &eof)) return NULL;
     GET_JOB_FROM_CAPSULE;
     if (eof == -1) eof = input_data_size > 0 ? 0 : 1;
@@ -70,7 +78,7 @@ iter_job(PyObject *self UNUSED, PyObject *args) {
         }
         if (result == RS_BLOCKED) break;
         Py_DECREF(ans);
-        PyErr_SetString(PyExc_RuntimeError, rs_strerror(result));
+        PyErr_SetString(RsyncError, rs_strerror(result));
         return NULL;
     }
     if ((ssize_t)output_size != PyBytes_GET_SIZE(ans)) {
@@ -79,14 +87,65 @@ iter_job(PyObject *self UNUSED, PyObject *args) {
     return Py_BuildValue("NO", ans, result == RS_DONE ? Py_True : Py_False);
 }
 
+static PyObject*
+begin_load_signature(PyObject *self UNUSED, PyObject *args UNUSED) {
+    rs_signature_t *sig = NULL;
+    rs_job_t *job = rs_loadsig_begin(&sig);
+    if (!sig || !job) return PyErr_NoMemory();
+    PyObject *jc = PyCapsule_New(job, JOB_CAPSULE, free_job_capsule);
+    if (!jc) { rs_job_free(job); rs_free_sumset(sig); return NULL; }
+    PyObject *sc = PyCapsule_New(sig, SIGNATURE_CAPSULE, free_sig_capsule);
+    if (!sc) { Py_DECREF(jc); rs_free_sumset(sig); return NULL; }
+    return Py_BuildValue("NN", jc, sc);
+}
+
+#define GET_SIG_FROM_CAPSULE \
+    rs_signature_t *sig = PyCapsule_GetPointer(sig_capsule, SIGNATURE_CAPSULE); \
+    if (!sig) { PyErr_SetString(PyExc_TypeError, "Not a sig capsule"); return NULL; }
+
+
+static PyObject*
+build_hash_table(PyObject *self UNUSED, PyObject *args) {
+    PyObject *sig_capsule;
+    if (!PyArg_ParseTuple(args, "O!y#|p", &PyCapsule_Type, &sig_capsule)) return NULL;
+    GET_SIG_FROM_CAPSULE;
+    rs_result res = rs_build_hash_table(sig);
+    if (res != RS_DONE) {
+        PyErr_SetString(RsyncError, rs_strerror(res));
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
+static PyObject*
+begin_create_delta(PyObject *self UNUSED, PyObject *args) {
+    PyObject *sig_capsule;
+    if (!PyArg_ParseTuple(args, "O!y#|p", &PyCapsule_Type, &sig_capsule)) return NULL;
+    GET_SIG_FROM_CAPSULE;
+    rs_job_t *job = rs_delta_begin(sig);
+    if (!job) return PyErr_NoMemory();
+    PyObject *ans = PyCapsule_New(job, JOB_CAPSULE, free_job_capsule);
+    if (!ans) rs_job_free(job);
+    return ans;
+}
+
+
 static PyMethodDef module_methods[] = {
     {"begin_create_signature", (PyCFunction)begin_create_signature, METH_VARARGS, ""},
+    {"begin_load_signature", (PyCFunction)begin_load_signature, METH_NOARGS, ""},
+    {"build_hash_table", (PyCFunction)build_hash_table, METH_VARARGS, ""},
+    {"begin_create_delta", (PyCFunction)begin_create_delta, METH_VARARGS, ""},
     {"iter_job", (PyCFunction)iter_job, METH_VARARGS, ""},
     {NULL, NULL, 0, NULL}        /* Sentinel */
 };
 
 static int
 exec_module(PyObject *m) {
+    RsyncError = PyErr_NewException("rsync.RsyncError", NULL, NULL);
+    if (RsyncError == NULL) return -1;
+    PyModule_AddObject(m, "RsyncError", RsyncError);
+
+
     PyModule_AddIntMacro(m, IO_BUFFER_SIZE);
     return 0;
 }
diff --git a/kittens/transfer/rsync.pyi b/kittens/transfer/rsync.pyi
index d9ce4b4a2..e9d904db7 100644
--- a/kittens/transfer/rsync.pyi
+++ b/kittens/transfer/rsync.pyi
@@ -7,9 +7,25 @@ class JobCapsule:
     pass
 
 
+class SignatureCapsule:
+    pass
+
+
+class RsyncError(Exception):
+    pass
+
+
 def begin_create_signature(file_size: int = -1, strong_len: int = 0) -> JobCapsule:
     pass
 
 
+def begin_load_signature() -> Tuple[JobCapsule, SignatureCapsule]:
+    pass
+
+
+def make_hash_table(sig: SignatureCapsule) -> None:
+    pass
+
+
 def iter_job(job_capsule: JobCapsule, input_data: bytes, eof: Optional[bool] = None) -> Tuple[bytes, bool]:
     pass