On-chain jetton processing

tolk 1.2;

struct Storage {
    owner: address;
    trustedJettonWallet: address?;
    jettonMinter: address;
}

fun Storage.load(): Storage {
    return Storage.fromCell(contract.getData());
}

fun Storage.save(self) {
    return contract.setData(self.toCell())
}

struct (0x7362d09c) JettonNotify {
    queryId: uint64
    jettonAmount: coins
    transferInitiator: address
    jettonVaultPayload: RemainingBitsAndRefs
}

struct (0x12345678) SetTrustedJettonWallet {
    trustedJettonWallet: address
}

type AllowedMessage = JettonNotify | SetTrustedJettonWallet;

fun onInternalMessage(in: InMessage) {
    val msg = lazy AllowedMessage.fromSlice(in.body);
    match (msg) {
        SetTrustedJettonWallet => {
            var storage = Storage.load();
            if (storage.owner != in.senderAddress || storage.trustedJettonWallet != null) {
                return;
            }
            storage.trustedJettonWallet = msg.trustedJettonWallet;
            storage.save();
        },
        JettonNotify => {
            var storage = Storage.load();
            if (storage.trustedJettonWallet == null || storage.trustedJettonWallet != in.senderAddress) {
                return;
            }
            // Process jettons
        }
    }
}

tolk 1.2

struct Storage {
    jettonMinterAddress: address
    jettonWalletAddress: address?
}

fun Storage.load(): Storage {
    return Storage.fromCell(contract.getData());
}

fun Storage.save(self) {
    return contract.setData(self.toCell())
}

struct (0x2c76b973) ProvideWalletAddress {
    queryId: uint64
    ownerAddress: address
    includeAddress: bool
}

struct (0xd1735400) TakeWalletAddress {
    queryId: uint64
    walletAddress: address
    ownerAddress: Cell<address>?
}

struct (0x7362d09c) JettonNotify {
    queryId: uint64
    jettonAmount: coins
    transferInitiator: address
    jettonVaultPayload: RemainingBitsAndRefs
}

// Unlike other opcodes here, this can be set arbitrarily
struct (0x12345678) InitContract {}

type AllowedMessage = InitContract | TakeWalletAddress | JettonNotify

fun onInternalMessage(in: InMessage) {
    val msg = lazy AllowedMessage.fromSlice(in.body);
    match (msg) {
        InitContract => {
            var storage = Storage.load();
            val msg = createMessage({
                bounce: true,
                value: 0,
                dest: storage.jettonMinterAddress,
                body: ProvideWalletAddress {
                    queryId: 0,
                    ownerAddress: contract.getAddress(),
                    includeAddress: false,
                },
            });
            msg.send(SEND_MODE_CARRY_ALL_REMAINING_MESSAGE_VALUE);
        }
        TakeWalletAddress => {
            var storage = Storage.load();
            if (storage.jettonMinterAddress != in.senderAddress) {
                return;
            }
            storage.jettonWalletAddress = msg.walletAddress;
            storage.save();
            // After receiving such message contract is initialized
        }
        JettonNotify => {
            var storage = Storage.load();
            if (storage.jettonWalletAddress == null || storage.jettonWalletAddress != in.senderAddress) {
                return;
            }
            // Process jettons
        }
    }
}

fun calculateJettonWallet(owner: address, jettonData: cell, jettonCode: cell, jettonMinter: address): address asm """
    c7 PUSHCTR
    // Unpack environment information from c7
    // https://docs.ton.org/tvm/registers#c7-%E2%80%94-environment-information-and-global-variables
    0 INDEX
    SWAP
    8 SETINDEX
    SWAP
    DUP
    ROTREV
    10 SETINDEX
    // Make it singleton
    1 TUPLE
    // owner md mc c7
    ROTREV
    CTOS            // owner_addr c7 md mc"
    2 PUSHINT       // owner_addr c7 md mc args"
    103289 PUSHINT  // owner_addr c7 md mc args get_jwa_method_id"
    5 0 REVERSE     // owner_addr get_jwa_method_id args mc md c7"
    // Moves the top stack value (ONE) to the third position: [A, B, C] -> [B, C, A]. We expect only 1 return value. Flag +256 for runvm enables this argument
    ONE 2 -ROLL
    // We use only these modes:
    // +1 = same_c3 (set c3 to code)
    // +4 = load c4 (persistent data) from stack and return its final value
    // +16 = load c7 (smart-contract context)
    // +256 = pop number N, return exactly N values from stack (only if res=0 or 1; if not enough then res=stk_und)
    // Mode 256 is crucial, because it ignores all stack values except the first one, and protects from stack poisoning
    277 RUNVM        // address exit_code c4'"
    2 BLKDROP       // address";
"""