mergesort in circom

mergesort.circom

// SPDX-License-Identifier: Apache-2.0
pragma circom 2.0.8;

// Merge sort - another variant that uses the grand product check for multiset equality,
// but requires a random seed

include "circomlib/comparators.circom";

/// @notice Merge pre-sorted left array and pre-sorted right array.
/// @param N     number of elements
/// @param bits  max number of bits in which each value in the list can be represented
/// @param asc   boolean value, 1 for ascending or 0 for descending direction of sort
template Merge(N, bits, asc) {
    assert(asc == 0 || asc == 1);

    var M = N \ 2;
    var R = M + N % 2;
    signal input randomSeed;
    signal input left[M];
    signal input right[R];
    signal output merged[N];

    // Merge computation (no constraints)
    var l = 0;
    var r = 0;
    var interm[N];
    for (var i = 0; i < N; i++) {
        var a = left[l];
        var b = right[r];
        var condition = asc == 1 ? a <= b : a > b; // asc or desc
        if ((condition && l < M) || r >= R) {
            interm[i] = left[l];
            l++;
        } else {
            interm[i] = right[r];
            r++;
        }
    }
    for (var i = 0; i < N; i++) {
        merged[i] <-- interm[i];
    }

    // Assert merge was actually computed correctly
    //  1. List is sorted
    component lt[N];
    for (var i = 1; i < N; i++) {
        lt[i] = LessEqThan(bits);
        lt[i].in[0] <== asc == 1 ? merged[i-1] : merged[i];
        lt[i].in[1] <== asc == 1 ? merged[i] : merged[i-1];
        lt[i].out === 1;
    }
    //  2. Check that merged list is multiset equal to {left || right} by using the grand product check
    //      ∏(a_i + γ) ?= ∏(b_i + γ)
    //      where γ is the random seed
    //      Ref:
    //          https://hackmd.io/@arielg/ByFgSDA7D
    signal grandProductInput[N];
    signal grandProductMerged[N];
    for (var i = 0; i < N; i++) {
        if (i < M) {
            grandProductInput[i] <== i > 0 ? grandProductInput[i-1] * (left[i] + randomSeed) : left[i] + randomSeed;
        } else {
            grandProductInput[i] <== grandProductInput[i-1] * (right[i-M] + randomSeed);
        }
        grandProductMerged[i] <== i > 0 ? grandProductMerged[i-1] * (merged[i] + randomSeed) : merged[i] + randomSeed;
    }
    grandProductInput[N-1] === grandProductMerged[N-1];
}

/// @notice Sort N elements by value. Values must be in the domain of 2^bits.
/// @param N     number of elements
/// @param bits  max number of bits in which each value in the list can be represented
/// @param asc   boolean value, 1 for ascending or 0 for descending direction of sort
template MergeSort(N, bits, asc) {
    assert(asc == 0 || asc == 1);
    signal input randomSeed;
    signal input unsorted[N];
    signal output sorted[N];

    // Random seed must be nonzero for the grand product check
    component isz = IsZero();
    isz.in <== randomSeed;
    isz.out === 0;

    component sortLeft;
    component sortRight;
    component merger;
    component midLt;
    signal intermUnsorted[2];

    if (N == 1) {
        sorted[0] <== unsorted[0];
    } else if (N == 2) {
        // Case base
        midLt = LessEqThan(bits);
        midLt.in[0] <== asc == 1 ? unsorted[0] : unsorted[1];
        midLt.in[1] <== asc == 1 ? unsorted[1] : unsorted[0];
        intermUnsorted[0] <== (1-midLt.out) * unsorted[1];
        intermUnsorted[1] <== (1-midLt.out) * unsorted[0];
        sorted[0] <== midLt.out * unsorted[0] + intermUnsorted[0];
        sorted[1] <== midLt.out * unsorted[1] + intermUnsorted[1];
    } else {
        var M = N \ 2;
        var R = M + N % 2;
        sortLeft = MergeSort(M, bits, asc);
        sortRight = MergeSort(R, bits, asc);
        sortLeft.randomSeed <== randomSeed;
        sortRight.randomSeed <== randomSeed;
        for (var i = 0; i < M; i++) {
            sortLeft.unsorted[i] <== unsorted[i];
        }
        for (var i = M; i < N; i++) {
            sortRight.unsorted[i-M] <== unsorted[i];
        }
        merger = Merge(N, bits, asc);
        merger.randomSeed <== randomSeed;
        for (var i = 0; i < M; i++) {
            merger.left[i] <== sortLeft.sorted[i];
        }
        for (var i = M; i < N; i++) {
            merger.right[i-M] <== sortRight.sorted[i-M];
        }
        for (var i = 0; i < N; i++) {
            sorted[i] <== merger.merged[i];
        }
    }
}

mergesort_main.circom

// SPDX-License-Identifier: Apache-2.0
pragma circom 2.0.8;

include "./mergesort.circom";

component main { public [ unsorted ] } = MergeSort(512, 64, 1);

/* INPUT = {
    "randomSeed": "694206942069420",
    "unsorted": [32,1,3,16,0,34,33,40,2,41,7,8,42,4,35,10,5,18,9,17,11,20,37,6,36,12,14,13,38,43,19,39,15,21,23,22,24,25,26,44,27,28,29,47,30,45,31,46,48,50,49,51,54,52,53,55,56,57,58,59,61,60,62,63,89,68,64,65,80,69,66,67,70,88,71,73,77,82,74,81,83,76,84,75,85,72,78,86,79,87,92,90,91,93,94,95,104,96,116,106,97,99,105,98,100,101,107,103,102,115,117,114,118,109,110,108,111,112,119,113,120,121,123,122,126,124,125,127,128,129,130,162,131,132,133,134,164,160,163,135,136,137,138,139,141,140,165,161,142,143,144,166,167,145,146,147,169,168,170,148,149,150,151,171,173,172,152,154,153,174,155,157,156,158,175,187,159,186,177,184,176,178,179,180,182,185,181,189,183,188,191,190,241,225,192,224,196,193,199,195,228,197,194,198,200,202,201,203,205,206,226,227,204,207,211,208,229,209,240,231,210,230,212,232,213,233,242,215,243,214,217,234,235,216,236,218,219,237,220,246,222,221,245,223,247,238,244,239,248,251,249,250,253,254,252,255,262,259,273,272,289,274,257,305,304,277,256,275,258,288,260,261,290,276,291,263,264,296,278,265,266,267,292,269,279,284,306,268,280,270,285,281,282,283,286,294,271,287,295,297,298,293,299,301,300,302,307,309,308,310,303,311,312,313,314,315,316,318,317,319,323,328,325,329,326,330,320,324,322,327,321,331,332,333,334,335,337,336,345,346,338,339,344,343,340,341,342,347,351,350,348,349,352,353,355,354,356,357,358,359,360,361,366,362,363,365,364,367,369,368,370,371,372,374,373,375,377,376,378,379,383,381,380,382,448,449,400,401,424,450,451,425,384,385,386,452,453,387,454,455,389,402,456,458,459,390,457,460,388,462,463,461,391,426,393,468,392,416,427,430,395,417,394,396,418,469,419,398,399,470,471,403,428,464,466,467,421,465,397,429,404,431,405,477,474,420,406,423,472,407,422,410,473,475,476,408,478,411,479,481,504,496,409,505,489,434,443,480,414,482,413,412,483,415,435,433,490,441,432,488,436,491,506,484,498,486,485,438,439,440,437,499,442,444,446,445,447,487,495,507,492,493,494,508,497,500,501,502,503,509,510,511]
} */