InstructionSet Zx_keccak {
    architectural_state {
        unsigned int XLEN=32;
        register unsigned<XLEN> X[32] [[is_main_reg]];
        extern unsigned<8> MEM[1 << XLEN] [[is_main_mem]];

        const unsigned int RV_CAUSE_ILLEGAL_INSTRUCTION = 0x02;

        register unsigned<64> K[32];
        register unsigned<8> LFSR;
    }
    functions {

        unsigned<64> ROL64(unsigned<64> val, unsigned<64> shift) {
            return (val << shift) | (val >> (64 - shift));
        }
        
        unsigned<64> readLane(unsigned int x, unsigned int y) {
            return K[x + 5 * y];
        }
        
        unsigned<64> XORLane(unsigned int x, unsigned int y, unsigned<64> value) {
            K[x + 5 * y] = K[x + 5 * y] ^ value;
            return 0;
        }

        unsigned<64> writeLane(unsigned int x, unsigned int y, unsigned<64> value) {
            K[x + 5 * y] = value;
            return 0;
        }

        unsigned<1> LFSR86540() {
            unsigned<1> result = (unsigned<1>)(LFSR & 0x01);
            if ((LFSR & 0x80) != 0) {
                // Primitive polynomial over GF(2): x^8+x^6+x^5+x^4+1
                LFSR = (LFSR << 1) ^ 0x71;
            } else {
                LFSR = LFSR << 1;
            }
            return result;
        }
    }
    instructions {
        // Custom-0 Opcode
        LK64 {
            encoding: imm[11:0] :: rs1[4:0] :: 3'b000 :: rd[4:0] :: 7'b0001011;
            assembly: "{name(rd)}, {imm}({name(rs1)})";
            behavior: {
                unsigned<XLEN> load_address = (unsigned<XLEN>)(X[rs1] + (signed<12>)imm);
                K[rd]= MEM[load_address+7:load_address];
            }
        }
        LK128 {
            encoding: imm[11:0] :: rs1[4:0] :: 3'b010 :: rd[4:0] :: 7'b0001011;
            assembly: "{name(rd)}, {imm}({name(rs1)})";
            behavior: if(rd%2 == 1) ; else {
                unsigned<XLEN> base_address = (unsigned<XLEN>)(X[rs1] + (signed<12>)imm);
                for(unsigned int i = 0; i<2; i++){
                    unsigned<XLEN> load_address = (unsigned<XLEN>)(base_address + 8*i);
                    K[rd+i]= MEM[load_address+7:load_address];
                }
            }
        }
        LK256 {
            encoding: imm[11:0] :: rs1[4:0] :: 3'b100 :: rd[4:0] :: 7'b0001011;
            assembly: "{name(rd)}, {imm}({name(rs1)})";
            behavior: if(rd%4 == 1) ; else {
                unsigned<XLEN> base_address = (unsigned<XLEN>)(X[rs1] + (signed<12>)imm);
                for(unsigned int i = 0; i<4; i++){
                    unsigned<XLEN> load_address = (unsigned<XLEN>)(base_address + 8*i);
                    K[rd+i]= MEM[load_address+7:load_address];
                }
            }
        }
        LK512 {
            encoding: imm[11:0] :: rs1[4:0] :: 3'b110 :: rd[4:0] :: 7'b0001011;
            assembly: "{name(rd)}, {imm}({name(rs1)})";
            behavior: if(rd%8 == 1) ; else {
                unsigned<XLEN> base_address = (unsigned<XLEN>)(X[rs1] + (signed<12>)imm);
                for(unsigned int i = 0; i<8; i++){
                    unsigned<XLEN> load_address = (unsigned<XLEN>)(base_address + 8*i);
                    K[rd+i]= MEM[load_address+7:load_address];
                }
            }
        }
        SK64 {
            encoding: imm[11:5] :: rs2[4:0] :: rs1[4:0] :: 3'b001 :: imm[4:0] :: 7'b0001011;
            assembly: "{name(rs2)}, {imm}({name(rs1)})";
            behavior: {
                unsigned<XLEN> store_address = (unsigned<XLEN>)(X[rs1] + (signed)imm);
                MEM[store_address+7:store_address] = K[rs2];
            }
        }
        SK128 {
            encoding: imm[11:5] :: rs2[4:0] :: rs1[4:0] :: 3'b011 :: imm[4:0] :: 7'b0001011;
            assembly: "{name(rs2)}, {imm}({name(rs1)})";
            behavior: if(rs2%2 == 1) ; else {
                unsigned<XLEN> base_address = (unsigned<XLEN>)(X[rs1] + (signed<12>)imm);
                for(unsigned int i = 0; i<2; i++){
                    unsigned<XLEN> store_address = (unsigned<XLEN>)(base_address + 8*i);
                    MEM[store_address+7:store_address] = K[rs2+i];
                }
            }
        }
        SK256 {
            encoding: imm[11:5] :: rs2[4:0] :: rs1[4:0] :: 3'b101 :: imm[4:0] :: 7'b0001011;
            assembly: "{name(rs2)}, {imm}({name(rs1)})";
            behavior: if(rs2%4 == 1) ; else {
                unsigned<XLEN> base_address = (unsigned<XLEN>)(X[rs1] + (signed<12>)imm);
                for(unsigned int i = 0; i<4; i++){
                    unsigned<XLEN> store_address = (unsigned<XLEN>)(base_address + 8*i);
                    MEM[store_address+7:store_address] = K[rs2+i];
                }
            }
        }
        SK512 {
            encoding: imm[11:5] :: rs2[4:0] :: rs1[4:0] :: 3'b111 :: imm[4:0] :: 7'b0001011;
            assembly: "{name(rs2)}, {imm}({name(rs1)})";
            behavior: if(rs2%8 == 1) ; else {
                unsigned<XLEN> base_address = (unsigned<XLEN>)(X[rs1] + (signed<12>)imm);
                for(unsigned int i = 0; i<8; i++){
                    unsigned<XLEN> store_address = (unsigned<XLEN>)(base_address + 8*i);
                    MEM[store_address+7:store_address] = K[rs2+i];
                }
            }
        }
        KECCAK_THETA {
            encoding: 7'b0000000 :: 5'b00000 :: 5'b00000 :: 3'b000 :: 5'b00000 :: 7'b0001011;
            assembly: "keccak.theta";
            behavior: {
                // Temporary storage for C array (5 columns)
                unsigned<64> C0, C1, C2, C3, C4;
                unsigned<64> D;
                
                // Compute the parity of the columns
                C0 = readLane(0, 0) ^ readLane(0, 1) ^ readLane(0, 2) ^ readLane(0, 3) ^ readLane(0, 4);
                C1 = readLane(1, 0) ^ readLane(1, 1) ^ readLane(1, 2) ^ readLane(1, 3) ^ readLane(1, 4);
                C2 = readLane(2, 0) ^ readLane(2, 1) ^ readLane(2, 2) ^ readLane(2, 3) ^ readLane(2, 4);
                C3 = readLane(3, 0) ^ readLane(3, 1) ^ readLane(3, 2) ^ readLane(3, 3) ^ readLane(3, 4);
                C4 = readLane(4, 0) ^ readLane(4, 1) ^ readLane(4, 2) ^ readLane(4, 3) ^ readLane(4, 4);
                
                // Apply theta effect to column 0
                D = C4 ^ ROL64(C1, 1);
                XORLane(0, 0, D);
                XORLane(0, 1, D);
                XORLane(0, 2, D);
                XORLane(0, 3, D);
                XORLane(0, 4, D);
                
                // Apply theta effect to column 1
                D = C0 ^ ROL64(C2, 1);
                XORLane(1, 0, D);
                XORLane(1, 1, D);
                XORLane(1, 2, D);
                XORLane(1, 3, D);
                XORLane(1, 4, D);
                
                // Apply theta effect to column 2
                D = C1 ^ ROL64(C3, 1);
                XORLane(2, 0, D);
                XORLane(2, 1, D);
                XORLane(2, 2, D);
                XORLane(2, 3, D);
                XORLane(2, 4, D);
                
                // Apply theta effect to column 3
                D = C2 ^ ROL64(C4, 1);
                XORLane(3, 0, D);
                XORLane(3, 1, D);
                XORLane(3, 2, D);
                XORLane(3, 3, D);
                XORLane(3, 4, D);
                
                // Apply theta effect to column 4
                D = C3 ^ ROL64(C0, 1);
                XORLane(4, 0, D);
                XORLane(4, 1, D);
                XORLane(4, 2, D);
                XORLane(4, 3, D);
                XORLane(4, 4, D);
            }
        }
        KECCAK_RHO_PI {
            encoding: 7'b0000000 :: 5'b00000 :: 5'b00000 :: 3'b001 :: 5'b00000 :: 7'b0001011;
            assembly: "keccak.rhopi";
            behavior: {
                unsigned<64> current, temp;
                unsigned int x, y;
                unsigned int r;
                unsigned int Y;
                
                // Start at coordinates (1, 0)
                x = 1;
                y = 0;
                current = readLane(x, y);
                
                // Iterate over ((0 1)(2 3))^t * (1 0) for 0 ≤ t ≤ 23
                for(unsigned int t = 0; t < 24; t++) {
                    // Compute the rotation constant r = (t+1)(t+2)/2
                    r = ((t + 1) * (t + 2) / 2) % 64;
                    
                    // Compute ((0 1)(2 3)) * (x y)
                    Y = (2 * x + 3 * y) % 5;
                    x = y;
                    y = Y;
                    
                    // Swap current and state(x,y), and rotate
                    temp = readLane(x, y);
                    writeLane(x, y, ROL64(current, r));
                    current = temp;
                }
            }
        }
        KECCAK_CHI {
            encoding: 7'b0000000 :: 5'b00000 :: 5'b00000 :: 3'b010 :: 5'b00000 :: 7'b0001011;
            assembly: "keccak.chi";
            behavior: {
                unsigned<64> T0, T1, T2, T3, T4;
                unsigned int y;
                for(y = 0; y < 5; y++) {
                    /* Take a copy of the plane */
                    T0 = readLane(0, y);
                    T1 = readLane(1, y);
                    T2 = readLane(2, y);
                    T3 = readLane(3, y);
                    T4 = readLane(4, y);
                    /* Compute χ on the plane */
                    writeLane(0, y, T0 ^ (~T1 & T2));
                    writeLane(1, y, T1 ^ (~T2 & T3));
                    writeLane(2, y, T2 ^ (~T3 & T4));
                    writeLane(3, y, T3 ^ (~T4 & T0));
                    writeLane(4, y, T4 ^ (~T0 & T1));
                }
            }
        }
        KECCAK_IOTA {
            encoding: 7'b0000000 :: 5'b00000 :: 5'b00000 :: 3'b011 :: 5'b00000 :: 7'b0001011;
            assembly: "keccak.iota";
            behavior: {
                unsigned int j;
                for (j = 0; j < 7; j++) {
                    unsigned int bitPosition = (unsigned int)((1 << j) - 1);
                    if (LFSR86540()) {
                        XORLane(0, 0, (unsigned<64>)1 << bitPosition);
                    }
                }
            }
        }
        KECCAK_LFSR_RESET {
            encoding: 7'b0000000 :: 5'b00000 :: 5'b00000 :: 3'b100 :: 5'b00000 :: 7'b0001011;
            assembly: "keccak.lfsr.reset";
            behavior: {
                LFSR = 0x01;
            }
        }
    }
}