MODULE SQLiteAggregate;
(* Aggregate functions, auxdata and the remaining UDF value/result APIs. *)

FROM SYSTEM IMPORT ADDRESS, ADR;
FROM CStr IMPORT CString, Copy;
FROM SQLite IMPORT
  DbHandle, StmtHandle, ContextHandle, ValueHandle,
  SQLiteOk, SQLiteRow, SQLiteDone, SQLiteNull,
  SQLiteInteger, SQLiteText, SQLiteTooBig, SQLiteNoMem,
  SQLiteOpenReadWrite, SQLiteOpenCreate, SQLiteOpenMemory, SQLiteUtf8,
  sqlite3_open_v2, sqlite3_close, sqlite3_exec,
  sqlite3_prepare_v2, sqlite3_step, sqlite3_column_int,
  sqlite3_column_text, sqlite3_finalize,
  sqlite3_create_function_v2,
  sqlite3_value_type, sqlite3_value_int64, sqlite3_value_numeric_type,
  sqlite3_value_dup, sqlite3_value_free,
  sqlite3_aggregate_context, sqlite3_get_auxdata, sqlite3_set_auxdata,
  sqlite3_result_int, sqlite3_result_int64, sqlite3_result_null,
  sqlite3_result_error, sqlite3_result_value, sqlite3_result_zeroblob,
  sqlite3_result_zeroblob64,
  sqlite3_result_text64, sqlite3_result_blob64,
  sqlite3_result_error_toobig, sqlite3_result_error_nomem,
  sqlite3_free, SQLiteArrayAt, SQLiteTransient;

TYPE
  SumPtr = POINTER TO LONGINT;

VAR
  ExitCode : INTEGER;
  db : DbHandle;
  stmt : StmtHandle;
  rc : INTEGER;
  ok : BOOLEAN;
  bin, auxMarker : ARRAY [0..7] OF CHAR;

PROCEDURE SameText (a, b : ARRAY OF CHAR) : BOOLEAN;
VAR
  i : INTEGER;
BEGIN
  i := 0;
  WHILE (a[i] # CHR(0)) AND (a[i] = b[i]) DO i := i + 1 END;
  RETURN (a[i] = CHR(0)) AND (b[i] = CHR(0))
END SameText;

PROCEDURE ExecSQL (sql : ARRAY OF CHAR) : BOOLEAN;
VAR
  errmsg : ADDRESS;
  good : BOOLEAN;
BEGIN
  errmsg := NIL;
  good := sqlite3_exec(db, sql, NIL, NIL, errmsg) = SQLiteOk;
  IF errmsg # NIL THEN sqlite3_free(errmsg) END;
  RETURN good
END ExecSQL;

PROCEDURE QueryInt (sql : ARRAY OF CHAR; want : INTEGER) : BOOLEAN;
VAR
  s : StmtHandle;
  code, got : INTEGER;
  good : BOOLEAN;
BEGIN
  s := NIL; got := -777; good := FALSE;
  code := sqlite3_prepare_v2(db, sql, -1, s, NIL);
  IF (code = SQLiteOk) AND (s # NIL) THEN
    IF sqlite3_step(s) = SQLiteRow THEN
      got := sqlite3_column_int(s, 0);
      IF sqlite3_step(s) = SQLiteDone THEN good := TRUE END
    END;
    IF sqlite3_finalize(s) # SQLiteOk THEN good := FALSE END
  END;
  RETURN good AND (got = want)
END QueryInt;

PROCEDURE QueryText (sql : ARRAY OF CHAR; want : ARRAY OF CHAR) : BOOLEAN;
VAR
  s : StmtHandle;
  code : INTEGER;
  p : CString;
  buf : ARRAY [0..63] OF CHAR;
  good : BOOLEAN;
BEGIN
  s := NIL; good := FALSE;
  code := sqlite3_prepare_v2(db, sql, -1, s, NIL);
  IF (code = SQLiteOk) AND (s # NIL) THEN
    IF sqlite3_step(s) = SQLiteRow THEN
      p := sqlite3_column_text(s, 0);
      IF (p # NIL) AND (Copy(p, buf) > 0) THEN
        good := SameText(buf, want)
      END;
      IF sqlite3_step(s) # SQLiteDone THEN good := FALSE END
    END;
    IF sqlite3_finalize(s) # SQLiteOk THEN good := FALSE END
  END;
  RETURN good
END QueryText;

PROCEDURE QueryError (sql : ARRAY OF CHAR; want : INTEGER) : BOOLEAN;
VAR
  s : StmtHandle;
  code : INTEGER;
  good : BOOLEAN;
BEGIN
  s := NIL; good := FALSE;
  code := sqlite3_prepare_v2(db, sql, -1, s, NIL);
  IF (code = SQLiteOk) AND (s # NIL) THEN
    good := sqlite3_step(s) = want;
    rc := sqlite3_finalize(s)
  END;
  RETURN good
END QueryError;

PROCEDURE Register (name : ARRAY OF CHAR; n : INTEGER;
                    xFunc, xStep, xFinal : ADDRESS) : BOOLEAN;
BEGIN
  RETURN sqlite3_create_function_v2(db, name, n, SQLiteUtf8,
    NIL, xFunc, xStep, xFinal, NIL) = SQLiteOk
END Register;

PROCEDURE SumStep (ctx : ContextHandle; argc : INTEGER; argv : ADDRESS);
VAR
  p : SumPtr;
  v : ValueHandle;
BEGIN
  p := VAL(SumPtr, sqlite3_aggregate_context(ctx, 8));
  IF p # NIL THEN
    v := SQLiteArrayAt(argv, 0);
    IF sqlite3_value_type(v) # SQLiteNull THEN
      p^ := p^ + sqlite3_value_int64(v)
    END
  END
END SumStep;

PROCEDURE SumFinal (ctx : ContextHandle);
VAR
  p : SumPtr;
BEGIN
  p := VAL(SumPtr, sqlite3_aggregate_context(ctx, 0));
  IF p = NIL THEN
    sqlite3_result_null(ctx)
  ELSE
    sqlite3_result_int64(ctx, p^)
  END
END SumFinal;

PROCEDURE AuxFn (ctx : ContextHandle; argc : INTEGER; argv : ADDRESS);
VAR
  p, q : ADDRESS;
BEGIN
  p := ADR(auxMarker[0]);
  sqlite3_set_auxdata(ctx, 0, p, NIL);
  q := sqlite3_get_auxdata(ctx, 0);
  IF q = p THEN
    sqlite3_result_int(ctx, 1)
  ELSE
    sqlite3_result_int(ctx, 0)
  END
END AuxFn;

PROCEDURE DupFn (ctx : ContextHandle; argc : INTEGER; argv : ADDRESS);
VAR
  v, d : ValueHandle;
BEGIN
  IF argc = 1 THEN
    v := SQLiteArrayAt(argv, 0);
    d := sqlite3_value_dup(v);
    IF d = NIL THEN
      sqlite3_result_error(ctx, "dup failed", -1)
    ELSE
      sqlite3_result_value(ctx, d);
      sqlite3_value_free(d)
    END
  ELSE
    sqlite3_result_error(ctx, "m2_dup expects 1 argument", -1)
  END
END DupFn;

PROCEDURE NumTypeFn (ctx : ContextHandle; argc : INTEGER; argv : ADDRESS);
BEGIN
  IF argc = 1 THEN
    sqlite3_result_int(ctx,
      sqlite3_value_numeric_type(SQLiteArrayAt(argv, 0)))
  ELSE
    sqlite3_result_error(ctx, "m2_numtype expects 1 argument", -1)
  END
END NumTypeFn;

PROCEDURE Zero32Fn (ctx : ContextHandle; argc : INTEGER; argv : ADDRESS);
BEGIN
  sqlite3_result_zeroblob(ctx, 4)
END Zero32Fn;

PROCEDURE Zero64Fn (ctx : ContextHandle; argc : INTEGER; argv : ADDRESS);
VAR
  code : INTEGER;
BEGIN
  code := sqlite3_result_zeroblob64(ctx, 4);
  IF code # SQLiteOk THEN
    sqlite3_result_error(ctx, "zeroblob64 failed", -1)
  END
END Zero64Fn;

PROCEDURE Text64Fn (ctx : ContextHandle; argc : INTEGER; argv : ADDRESS);
BEGIN
  sqlite3_result_text64(ctx, "hi", 2, SQLiteTransient(), SQLiteUtf8)
END Text64Fn;

PROCEDURE Blob64Fn (ctx : ContextHandle; argc : INTEGER; argv : ADDRESS);
BEGIN
  sqlite3_result_blob64(ctx, ADR(bin[0]), 3, SQLiteTransient())
END Blob64Fn;

PROCEDURE TooBigFn (ctx : ContextHandle; argc : INTEGER; argv : ADDRESS);
BEGIN
  sqlite3_result_error_toobig(ctx)
END TooBigFn;

PROCEDURE NoMemFn (ctx : ContextHandle; argc : INTEGER; argv : ADDRESS);
BEGIN
  sqlite3_result_error_nomem(ctx)
END NoMemFn;

BEGIN
  ExitCode := 1;
  db := NIL; stmt := NIL; ok := FALSE;
  bin[0] := CHR(0); bin[1] := CHR(1); bin[2] := CHR(2);
  auxMarker[0] := 'A';

  IF (sqlite3_open_v2(":memory:", db,
        SQLiteOpenReadWrite + SQLiteOpenCreate + SQLiteOpenMemory,
        NIL) = SQLiteOk) AND (db # NIL)
     AND ExecSQL("CREATE TABLE nums(x INTEGER);")
     AND ExecSQL("INSERT INTO nums VALUES(1);")
     AND ExecSQL("INSERT INTO nums VALUES(2);")
     AND ExecSQL("INSERT INTO nums VALUES(3);") THEN
    ok := TRUE
  END;

  IF ok THEN
    ok := Register("m2_sum", 1, NIL, ADR(SumStep), ADR(SumFinal))
          AND QueryInt("SELECT m2_sum(x) FROM nums;", 6)
  END;
  IF ok THEN
    ok := Register("m2_aux", 0, ADR(AuxFn), NIL, NIL)
          AND QueryInt("SELECT m2_aux();", 1)
  END;
  IF ok THEN
    ok := Register("m2_dup", 1, ADR(DupFn), NIL, NIL)
          AND QueryText("SELECT m2_dup('abc');", "abc")
  END;
  IF ok THEN
    ok := Register("m2_numtype", 1, ADR(NumTypeFn), NIL, NIL)
          AND QueryInt("SELECT m2_numtype('42');", SQLiteInteger)
  END;
  IF ok THEN
    ok := Register("m2_zero32", 0, ADR(Zero32Fn), NIL, NIL)
          AND QueryText("SELECT hex(m2_zero32());", "00000000")
  END;
  IF ok THEN
    ok := Register("m2_zero64", 0, ADR(Zero64Fn), NIL, NIL)
          AND QueryText("SELECT hex(m2_zero64());", "00000000")
  END;
  IF ok THEN
    ok := Register("m2_text64", 0, ADR(Text64Fn), NIL, NIL)
          AND QueryText("SELECT m2_text64();", "hi")
  END;
  IF ok THEN
    ok := Register("m2_blob64", 0, ADR(Blob64Fn), NIL, NIL)
          AND QueryText("SELECT hex(m2_blob64());", "000102")
  END;
  IF ok THEN
    ok := Register("m2_toobig", 0, ADR(TooBigFn), NIL, NIL)
          AND QueryError("SELECT m2_toobig();", SQLiteTooBig)
  END;
  IF ok THEN
    ok := Register("m2_nomem", 0, ADR(NoMemFn), NIL, NIL)
          AND QueryError("SELECT m2_nomem();", SQLiteNoMem)
  END;

  IF db # NIL THEN
    IF sqlite3_close(db) # SQLiteOk THEN ok := FALSE END
  END;
  IF ok THEN ExitCode := 42 END
END SQLiteAggregate.