iden3 / ffwasm Goto Github PK

use wasmer_runtime::{
    Instance,
    Module,
    types::{
        MemoryDescriptor
    },
    units::*,
    memory::{
        Memory
    },
    Ctx,
    instantiate,
    DynFunc,
    Value,
    Func,
    imports,
    func,
    error,
};
use std::marker::Sync;

pub struct WasmInstance(Instance);

impl WasmInstance {
    const MEMORY_INIT:u32 = 32767;

    pub fn new(_bytes:&[u8]) -> Self {
        let descriptor = MemoryDescriptor::new(Pages(Self::MEMORY_INIT), None, false).unwrap();
        let mem = Memory::new(descriptor).unwrap();
        let import_object = imports! {
            "env"=>{
                "memory"=>mem
            },
            "runtime"=>{
                "error"=>func!(Self::error),
                "log"=>func!(Self::log),
                "logGetSignal"=>func!(Self::log_get_signal),
                "logSetSignal"=>func!(Self::log_set_signal),
                "logStartComponent"=>func!(Self::log_start_component),
                "logFinishComponent"=>func!(Self::log_finish_component),
            }
        };
        
        let mut instance = instantiate(_bytes, &import_object).unwrap();
        Self(instance)
    }

    fn error(_:&mut Ctx, code:i32, pstr:i32, a:i32, b:i32, c:i32, d:i32) {
        println!("Error in Wasm. code:{}, pstr:{}, a:{}, b:{}, c:{}, d:{}",code,pstr,a,b,c,d);
        panic!("Error in Wasm. code:{}, pstr:{}, a:{}, b:{}, c:{}, d:{}",code,pstr,a,b,c,d);
    }

    fn log(_:&mut Ctx,in0:i32) {

    }

    fn log_get_signal(_:&mut Ctx, signal:i32, p_val:i32) {
        println!("signal: {}",signal);
        println!("pVal: {}",p_val);
    }

    fn log_set_signal(_:&mut Ctx,in0:i32,in1:i32) {

    }

    fn log_start_component(_:&mut Ctx,in0:i32) {

    }

    fn log_finish_component(_:&mut Ctx,in0:i32) {

    }

    pub fn init(&self,inputs:[i32;1]) {
        let init_func:Func<i32,()> = self.0.exports.get("init").unwrap();
        init_func.call(inputs[0]).unwrap();
    }

    pub fn get_signal_offset32(&self,inputs:[i32;4]) {
        let get_sig_func:Func<(i32,i32,i32,i32),()> = self.0.exports.get("getSignalOffset32").unwrap();
        get_sig_func.call(inputs[0],inputs[1],inputs[2],inputs[3]).unwrap();
    }

    pub fn get_multi_signal(&self,inputs:[i32;5]) {
        let get_multi_sig_func:Func<(i32,i32,i32,i32,i32),()> = self.0.exports.get("multiGetSignal").unwrap();
        get_multi_sig_func.call(inputs[0],inputs[1],inputs[2],inputs[3],inputs[4]).unwrap();
    }

    pub fn set_signal(&self,inputs:[i32;4]) {
        let set_sig_func:Func<(i32,i32,i32,i32),()> = self.0.exports.get("setSignal").unwrap();
        set_sig_func.call(inputs[0],inputs[1],inputs[2],inputs[3]).unwrap();
    }

    pub fn get_p_witness(&self,inputs:[i32;1]) -> i32 {
        let get_pw_func:Func<i32,i32> = self.0.exports.get("getPWitness").unwrap();
        get_pw_func.call(inputs[0]).unwrap()
    }

    pub fn get_witness_buffer(&self) -> i32 {
        let get_pw_buf_func:Func<(),i32> = self.0.exports.get("getWitnessBuffer").unwrap();
        get_pw_buf_func.call().unwrap()
    }

    pub fn get_fr_len(&self) -> i32 {
        let get_fr_func:Func<(),i32> = self.0.exports.get("getFrLen").unwrap();
        get_fr_func.call().unwrap()
    }

    pub fn get_p_raw_prime(&self) -> i32 {
        let get_praw_func:Func<(),i32> = self.0.exports.get("getPRawPrime").unwrap();
        get_praw_func.call().unwrap()
    }

    pub fn get_n_vars(&self) -> i32 {
        let get_nvars_func:Func<(),i32> = self.0.exports.get("getNVars").unwrap();
        get_nvars_func.call().unwrap()
    }

    pub fn get_memory(&self) -> &Memory {
        self.0.context().memory(0)
    }

}

unsafe impl Sync for WasmInstance {}

use neon::prelude::*;
use crate::{WasmInstance};
use lazy_static::lazy_static;
use wasmer_runtime::{
    WasmPtr,
    Instance,
    Module,
    types::{
        MemoryDescriptor
    },
    units::*,
    memory::{
        Memory
    },
    Ctx,
    instantiate,
    DynFunc,
    Value,
    Func,
    imports,
    func,
    error,
};

lazy_static! {
    static ref INSTANCE: WasmInstance = {
        let bytes = include_bytes!("../test_circom/mul_bn128.wasm");
        WasmInstance::new(bytes)
    };
}


register_module!(mut cx, {
    cx.export_function("init", init);
    cx.export_function("getSignalOffset32", get_signal_offset32);
    cx.export_function("getMultiSignal", get_multi_signal);
    cx.export_function("setSignal", set_signal);
    cx.export_function("getPWitness", get_p_witness);
    cx.export_function("getWitnessBuffer", get_witness_buffer);
    cx.export_function("getFrLen", get_fr_len);
    cx.export_function("getPRawPrime", get_p_raw_prime);
    cx.export_function("getNVars", get_n_vars);
    cx.export_function("setU32Memory", set_u32_memory);
    cx.export_function("getU8Memory", get_u8_memory);
    cx.export_function("getU32Memory", get_u32_memory)
});


pub fn init(mut cx: FunctionContext) -> JsResult<JsUndefined> {
    const INPUT_SIZE:usize = 1;
    let mut inputs:[i32;INPUT_SIZE] = [0;INPUT_SIZE];
    for i in 0..INPUT_SIZE {
        inputs[i] = cx.argument::<JsNumber>(i as i32)?.value() as i32;
    }
    INSTANCE.init(inputs);
    Ok(cx.undefined())
}

pub fn get_signal_offset32(mut cx: FunctionContext) -> JsResult<JsUndefined> {
    const INPUT_SIZE:usize = 4;
    let mut inputs:[i32;INPUT_SIZE] = [0;INPUT_SIZE];
    for i in 0..INPUT_SIZE {
        inputs[i] = cx.argument::<JsNumber>(i as i32)?.value() as i32;
    }
    INSTANCE.get_signal_offset32(inputs);
    Ok(cx.undefined())
}

pub fn get_multi_signal(mut cx: FunctionContext) -> JsResult<JsUndefined> {
    const INPUT_SIZE:usize = 5;
    let mut inputs:[i32;INPUT_SIZE] = [0;INPUT_SIZE];
    for i in 0..INPUT_SIZE {
        inputs[i] = cx.argument::<JsNumber>(i as i32)?.value() as i32;
    }
    INSTANCE.get_multi_signal(inputs);
    Ok(cx.undefined())
}

pub fn set_signal(mut cx: FunctionContext) -> JsResult<JsUndefined> {
    const INPUT_SIZE:usize =4;
    let mut inputs:[i32;INPUT_SIZE] = [0;INPUT_SIZE];
    for i in 0..INPUT_SIZE {
        inputs[i] = cx.argument::<JsNumber>(i as i32)?.value() as i32;
    }
    INSTANCE.set_signal(inputs);
    Ok(cx.undefined())
}

pub fn get_p_witness(mut cx: FunctionContext) -> JsResult<JsNumber> {
    const INPUT_SIZE:usize = 1;
    let mut inputs:[i32;INPUT_SIZE] = [0;INPUT_SIZE];
    for i in 0..INPUT_SIZE {
        inputs[i] = cx.argument::<JsNumber>(i as i32)?.value() as i32;
    }
    let result = INSTANCE.get_p_witness(inputs);
    Ok(cx.number(result))
}

pub fn get_witness_buffer(mut cx: FunctionContext) -> JsResult<JsNumber> {
    let result = INSTANCE.get_witness_buffer();
    Ok(cx.number(result))
}

pub fn get_fr_len(mut cx: FunctionContext) -> JsResult<JsNumber> {
    let result = INSTANCE.get_fr_len();
    Ok(cx.number(result))
}

pub fn get_p_raw_prime(mut cx: FunctionContext) -> JsResult<JsNumber> {
    let result = INSTANCE.get_p_raw_prime();
    Ok(cx.number(result))
}

pub fn get_n_vars(mut cx: FunctionContext) -> JsResult<JsNumber> {
    let result = INSTANCE.get_n_vars();
    Ok(cx.number(result))
}


pub fn set_u32_memory(mut cx: FunctionContext) -> JsResult<JsUndefined> {
    let offset = cx.argument::<JsNumber>(0)?.value() as u32;
    let value = cx.argument::<JsNumber>(1)?.value() as u32;
    let v_bytes = value.to_le_bytes();
    let memory = INSTANCE.get_memory();
    for i in 0..4{
        let ptr = WasmPtr::<u8>::new(4*offset+i as u32);
        let derefed_ptr = ptr.deref(memory).unwrap();
        derefed_ptr.set(v_bytes[i as usize] as u8);
    }
    Ok(cx.undefined())
}

pub fn get_u8_memory(mut cx: FunctionContext) -> JsResult<JsNumber> {
    let offset = cx.argument::<JsNumber>(0)?.value() as u32;
    let value = _get_u8_memory(offset);
    Ok(cx.number(f64::from(value)))
}

pub fn get_u32_memory(mut cx: FunctionContext) -> JsResult<JsNumber> {
    let offset = cx.argument::<JsNumber>(0)?.value() as u32;
    let mut bytes4:[u8;4] = [0;4];
    for i in 0..4 {
        let value = _get_u8_memory(4*offset+i as u32);
        bytes4[i] = value;
    }
    let value:u32 = u32::from_le_bytes(bytes4);
    Ok(cx.number(f64::from(value)))
}

fn _get_u8_memory(offset:u32) -> u8 {
    let memory = INSTANCE.get_memory();
    let ptr = WasmPtr::<u8>::new(offset);
    //println!("u8 ptr is {:?}", ptr);
    let derefed_ptr = ptr.deref(memory).unwrap();
    let value:u8 = derefed_ptr.get();
    value
}

const utils = require("./utils");
const Scalar = require("ffjavascript").Scalar;
const F1Field = require("ffjavascript").F1Field;
const {ffi} = require("./index");
const assert = require('assert');

module.exports = async function builder(code, options) {

   options = options || {};

   const memory = new WebAssembly.Memory({initial:32767});
   const wasmModule = await WebAssembly.compile(code);
   let wc;

   const instance = await WebAssembly.instantiate(wasmModule, {
       env: {
           "memory": memory
       },
       runtime: {
           error: function(code, pstr, a,b,c,d) {
               let errStr;
               if (code == 7) {
                   errStr=p2str(pstr) + " " + wc.getFr(b).toString() + " != " + wc.getFr(c).toString() + " " +p2str(d);
               } else if (code == 9) {
                   errStr=p2str(pstr) + " " + wc.getFr(b).toString() + " " +p2str(c);
               } else if ((code == 5)&&(options.sym)) {
                   errStr=p2str(pstr)+ " " + options.sym.labelIdx2Name[c];
               } else {
                   errStr=p2str(pstr)+ " " + a + " " + b + " " + c + " " + d;
               }
               console.log("ERROR: ", code, errStr);
               throw new Error(errStr);
           },
           log: function(a) {
               //console.log(wc.getFr(a).toString());
           },
           logGetSignal: function(signal, pVal) {
               if (options.logGetSignal) {
                   options.logGetSignal(signal, wc.getFr(pVal) );
               }
           },
           logSetSignal: function(signal, pVal) {
               if (options.logSetSignal) {
                   options.logSetSignal(signal, wc.getFr(pVal) );
               }
           },
           logStartComponent: function(cIdx) {
               if (options.logStartComponent) {
                   options.logStartComponent(cIdx);
               }
           },
           logFinishComponent: function(cIdx) {
               if (options.logFinishComponent) {
                   options.logFinishComponent(cIdx);
               }
           }
       }
   });

   const sanityCheck =
       options &&
       (
           options.sanityCheck ||
           options.logGetSignal ||
           options.logSetSignal ||
           options.logStartComponent ||
           options.logFinishComponent
       );

   wc = new WitnessCalculator(memory, instance, sanityCheck);
   return wc;

   function p2str(p) {
       const i8 = new Uint8Array(ffi.memoryAsArrayBuffer());

       const bytes = [];

       for (let i=0; i8[p+i]>0; i++)  bytes.push(i8[p+i]);

       return String.fromCharCode.apply(null, bytes);
   }
};

class WitnessCalculator{
   constructor(memory, instance, sanityCheck) {
       this.memory = memory;
       this.i32 = new Uint32Array(memory.buffer);
       this.instance = instance;
       this.n32 = (ffi.getFrLen() >> 2) - 2;
       assert(this.n32===(this.instance.exports.getFrLen() >> 2) - 2);
       const pRawPrime = ffi.getPRawPrime();
       assert(pRawPrime===this.instance.exports.getPRawPrime());
       const arr = new Array(this.n32);
       for (let i=0; i<this.n32; i++) {
           arr[this.n32-1-i] = ffi.getU32Memory((pRawPrime >> 2) + i);
       }
       this.prime = Scalar.fromArray(arr, 0x100000000);
       this.Fr = new F1Field(this.prime);
       this.mask32 = Scalar.fromString("FFFFFFFF", 16);
       this.NVars = ffi.getNVars();
       this.n64 = Math.floor((this.Fr.bitLength - 1) / 64)+1;
       this.R = this.Fr.e( Scalar.shiftLeft(1 , this.n64*64));
       this.RInv = this.Fr.inv(this.R);
       this.sanityCheck = sanityCheck;
   }

   async _doCalculateWitness(input, sanityCheck) {
       ffi.init((this.sanityCheck || sanityCheck) ? 1 : 0);
       this.instance.exports.init((this.sanityCheck || sanityCheck) ? 1 : 0);
       const pSigOffset = this.allocInt();
       const pFr = this.allocFr();
       const keys = Object.keys(input);
       const n = keys.length;
       let i;
       keys.forEach( (k) => {
           const h = utils.fnvHash(k);
           const hMSB = parseInt(h.slice(0,8), 16);
           const hLSB = parseInt(h.slice(8,16), 16);
           /// comparing values of the wasm memory in the Rust environment with those of WebAssembly.Memory in js.
           for(i=0;i<32767;i++){
               try{
                   assert(this.i32[i]===ffi.getU32Memory(i));
               }
               catch(e){
                   console.log(e);
                   return;
               }
           }
           ///
           try {
               this.instance.exports.getSignalOffset32(pSigOffset, 0, hMSB, hLSB);
               ffi.getSignalOffset32(pSigOffset, 0, hMSB, hLSB);
           } catch (err) {
               console.error(err);
               throw new Error(`Signal ${k} is not an input of the circuit.`);
           }
           const sigOffset = this.getInt(pSigOffset);
           const fArr = utils.flatArray(input[k]);
           for (let i=0; i<fArr.length; i++) {
               this.setFr(pFr, fArr[i]);
               ffi.setSignal(0, 0, sigOffset + i, pFr);
               this.instance.exports.setSignal(0, 0, sigOffset + i, pFr);
           }
       });
   }

   async calculateWitness(input, sanityCheck) {
       const self = this;

       const old0 = ffi.getU32Memory(0);//self.i32[0];
       const w = [];

       await self._doCalculateWitness(input, sanityCheck);

       for (let i=0; i<self.NVars; i++) {
           const pWitness = ffi.getPWitness(i);
           w.push(self.getFr(pWitness));
       }
       
       ffi.setU32Memory(0,old0);
       return w;
   }

   async calculateBinWitness(input, sanityCheck) {
       const self = this;
       const old0 = ffi.getU32Memory(0);//self.i32[0];
       await self._doCalculateWitness(input, sanityCheck);

       const pWitnessBuffer = ffi.getWitnessBuffer();

       ffi.setU32Memory(0,old0);
       self.i32[0] = old0;
       const memory_buf = new Uint8Array(self.memory.buffer);
       const buff = self.sliceU8Memory(pWitnessBuffer,pWitnessBuffer + (self.NVars * self.n64 * 8));
       return buff;
   }

   sliceU8Memory(init, finish) {
       let i;
       let values = [];
       for(i=init; i<finish; i++){
           values.push(ffi.getU8Memory(i));
       }
       return new Uint8Array(values);
   }

   allocInt() {
       const p = ffi.getU32Memory(0);//this.i32[0];
       ffi.setU32Memory(0,p+8); this.i32[0] = this.i32[0]+8;
       return p;
   }

   allocFr() {
       const p = ffi.getU32Memory(0);//this.i32[0];
       this.i32[0] = this.i32[0]+this.n32*4 + 8;
       ffi.setU32Memory(0, p+this.n32*4 + 8);
       return p;
   }

   getInt(p) {
       return ffi.getU32Memory(p>>2);//this.i32[p>>2];
   }

   setInt(p, v) {
       ffi.setU32Memory(p>>2, v);
   }

   getFr(p) {
       const self = this;
       const idx = (p>>2);

       if (ffi.getU32Memory(idx+1) & 0x80000000) {
           const arr = new Array(self.n32);
           for (let i=0; i<self.n32; i++) {
               arr[self.n32-1-i] = ffi.getU32Memory(idx+2*i);
           }
           const res = self.Fr.e(Scalar.fromArray(arr, 0x100000000));
           if (ffi.getU32Memory(idx + 1) & 0x40000000) {
               return fromMontgomery(res);
           } else {
               return res;
           }

       } else {
           if (ffi.getU32Memory(idx) & 0x80000000) {
               return self.Fr.e( ffi.getU32Memory(idx) - 0x100000000);
           } else {
               return self.Fr.e(ffi.getU32Memory(idx));
           }
       }

       function fromMontgomery(n) {
           return self.Fr.mul(self.RInv, n);
       }

   }


   setFr(p, v) {
       const self = this;

       v = self.Fr.e(v);

       const minShort = self.Fr.neg(self.Fr.e("80000000", 16));
       const maxShort = self.Fr.e("7FFFFFFF", 16);

       if (  (self.Fr.geq(v, minShort))
           &&(self.Fr.leq(v, maxShort)))
       {
           let a;
           if (self.Fr.geq(v, self.Fr.zero)) {
               a = Scalar.toNumber(v);
           } else {
               a = Scalar.toNumber( self.Fr.sub(v, minShort));
               a = a - 0x80000000;
               a = 0x100000000 + a;
           }
           ffi.setU32Memory(p>>2,a);
           ffi.setU32Memory((p >> 2) + 1,0);
           self.i32[(p >> 2)] = a;
           self.i32[(p >> 2) + 1] = 0;
           return;
       }

       ffi.setU32Memory(p>>2,0);
       ffi.setU32Memory((p >> 2) + 1,0x80000000);
       self.i32[(p >> 2)] = 0;
       self.i32[(p >> 2) + 1] = 0x80000000;
       const arr = Scalar.toArray(v, 0x100000000);
       for (let i=0; i<self.n32; i++) {
           const idx = arr.length-1-i;

           if ( idx >=0) {
               ffi.setU32Memory((p >> 2) + 2 + i,arr[idx]);
               self.i32[(p >> 2) + 2 + i] = arr[idx];
           } else {
               ffi.setU32Memory((p >> 2) + 2 + i,0);
               self.i32[(p >> 2) + 2 + i] = 0;
           }
       }
   }
}

template Mul() {
    signal input in[2];
    signal output out;

    signal dbl <== in[1] * in[1];
    out <== in[0] * dbl;
}

component main = Mul();