4-2. HEaaN End to End Example

////////////////////////////////////////////////////////////////////////////////
//                                                                            //
// Copyright (C) 2021-2025 CryptoLab Inc.                                    //
//                                                                            //
// - This file is part of HEaaN homomorphic encryption library.               //
// - HEaaN cannot be copied and/or distributed without the express permission //
//  of CryptoLab Inc.                                                        //
//                                                                            //
////////////////////////////////////////////////////////////////////////////////

#include "HEaaN/HEaaN.hpp"
#include <iostream>
#include <random>
/*
In this example, we describe performing HEaaN on GPU. The following code
performs the same operations as the code in '5-bootstrap.cpp'.
*/

inline void printMessage(const HEaaN::Message &msg, bool is_complex = true,
                         size_t start_num = 2, size_t end_num = 2,
                         const int precision = 10) {
    const size_t msg_size = msg.getSize();
    std::cout.precision(precision);
    std::cout << "[ ";
    for (size_t i = 0; i < start_num; ++i) {
        if (is_complex)
            std::cout << msg[i] << ", ";
        else
            std::cout << msg[i].real() << ", ";
    }
    std::cout << "..., ";
    for (size_t i = end_num; i > 1; --i) {
        if (is_complex)
            std::cout << msg[msg_size - i] << ", ";
        else
            std::cout << msg[msg_size - i].real() << ", ";
    }
    if (is_complex)
        std::cout << msg[msg_size - 1] << " ]" << std::endl;
    else
        std::cout << msg[msg_size - 1].real() << " ]" << std::endl;
}

int main() {
    if (!HEaaN::CudaTools::isAvailable()) {
        std::cout << "Device is not available." << std::endl;
        return -1;
    }

    // You can use other bootstrappable parameter instead of FGb.
    // See 'include/HEaaN/ParameterPreset.hpp' for more details.
    HEaaN::ParameterPreset preset = HEaaN::ParameterPreset::FGb;
    HEaaN::Context context = makeContext(preset, {0}); // set CUDA device ID

    const auto log_slots = getLogFullSlots(context);

    HEaaN::SecretKey sk(context);
    HEaaN::KeyPack pack(context);
    HEaaN::KeyGenerator keygen(context, sk, pack);

    std::cout << "Generate encryption key ... " << std::endl;
    keygen.genEncKey();
    std::cout << "done" << std::endl << std::endl;

    std::cout << "Generate multiplication key and conjugation key ... "
              << std::endl;
    keygen.genMultKey();
    keygen.genConjKey();
    std::cout << "done" << std::endl << std::endl;

    std::cout << "Generate rotation keys used in the bootstrap process ..."
              << std::endl;
    keygen.genRotKeysForBootstrap(log_slots);
    std::cout << "done" << std::endl << std::endl;

    std::cout << "Load all public keys to GPU memory ..." << std::endl;
    pack.to(HEaaN::getCurrentCudaDevice());
    std::cout << std::endl;

    HEaaN::Encryptor enc(context);
    HEaaN::Decryptor dec(context);
    HEaaN::HomEvaluator eval(context, pack);

    /*
    Bootstrapper constructor pre-compute the constants for bootstrapping.
    */
    std::cout << "Generate Bootstrapper (including pre-computing constants for "
                 "bootstrapping) ..."
              << std::endl;
    HEaaN::Bootstrapper btp(eval);
    std::cout << std::endl;

    std::cout << "Load boot constants data to GPU memory ..." << std::endl;
    btp.loadBootConstants(log_slots, HEaaN::getCurrentCudaDevice());

    {
        static std::default_random_engine gen{std::random_device()()};
        std::uniform_real_distribution<double> dist(-1.0L, 1.0L);
        
        HEaaN::Message msg(log_slots);
        for (size_t i = 0; i < msg.getSize(); ++i) {
            msg[i].real(dist(gen));
            msg[i].imag(dist(gen));
        }

        std::cout << std::endl << "Input message : " << std::endl;
        printMessage(msg);
        std::cout << std::endl;

        /*
        Use 'to' function for loading object to GPU memory.
        For all HEaaN function, if input data is on GPU then output data is also
        on GPU.
        */
        msg.to(HEaaN::getCurrentCudaDevice());

        std::cout << "Encrypt ... ";
        HEaaN::Ciphertext ctxt(context);
        enc.encrypt(msg, pack, ctxt);
        std::cout << "done" << std::endl << std::endl;

        std::cout << "Input ciphertext - level " << ctxt.getLevel()
                  << std::endl;

        // consume level with computation
        std::cout << std::endl
                  << "Computing the polynomial 0.1*X^4 - 0.1X + 0.4 ... ";
        HEaaN::Ciphertext ctxt2(context), ctxt4(context), ctxt_rst(context);
        eval.mult(ctxt, ctxt, ctxt2);
        eval.mult(ctxt2, ctxt2, ctxt4);
        eval.mult(ctxt4, 0.1, ctxt_rst);
        eval.mult(ctxt, 0.1, ctxt);
        eval.sub(ctxt_rst, ctxt, ctxt_rst);
        eval.add(ctxt_rst, 0.4, ctxt_rst);
        std::cout << "done" << std::endl;

        std::cout << "Result ciphertext - level " << ctxt_rst.getLevel()
                  << std::endl
                  << std::endl;

        HEaaN::Ciphertext ctxt_boot(context);
        std::cout << "Bootstrapping ... " << std::endl;
        /*
        This is the bootstrapping for complex numbers.
        You can omit the third argument if encrypted messages are all real
        numbers. The bootstrapping for real numbers is 30~40% faster.
        */
        btp.bootstrap(ctxt_rst, ctxt_boot, true);
        
        std::cout << std::endl;

        std::cout << "Result ciphertext after bootstrapping - level "
                  << ctxt_boot.getLevel() << std::endl
                  << std::endl;

        /*
        Decrypt does not work well
        if ciphertext and key data are on different device memory.
        */
        sk.to(HEaaN::getCurrentCudaDevice());

        HEaaN::Message dmsg, dmsg_boot;
        std::cout << "Decrypt ... ";
        dec.decrypt(ctxt_rst, sk, dmsg);
        dec.decrypt(ctxt_boot, sk, dmsg_boot);
        std::cout << "done" << std::endl;

        /*
        Decrypted message data should be on CPU memory
        because below code is not working for GPU
        */
        dmsg.to(HEaaN::getDefaultDevice());
        dmsg_boot.to(HEaaN::getDefaultDevice());

        std::cout.precision(10);
        std::cout << std::endl << "Decrypted message : " << std::endl;
        printMessage(dmsg);
        std::cout << std::endl
                  << "Decrypted message (after bootstrapping) : " << std::endl;
        printMessage(dmsg_boot);
    }

    return 0;
}