SimpleStorage.java

import com.google.common.base.Preconditions;
import net.fabricmc.fabric.api.transfer.v1.storage.Storage;
import net.fabricmc.fabric.api.transfer.v1.storage.StorageView;
import net.fabricmc.fabric.api.transfer.v1.transaction.Participant;
import net.fabricmc.fabric.api.transfer.v1.transaction.Transaction;
import net.fabricmc.fabric.api.transfer.v1.transaction.TransactionResult;
import org.jetbrains.annotations.Nullable;

public class SimpleStorage<T> implements Storage<T>, StorageView<T>, Participant<Long> {
    private final long maxAmount;
    private long amount;
    private T resource;
    private int version;

    public SimpleStorage(long maxAmount) {
        this(maxAmount, 0, null);
    }

    public SimpleStorage(long maxAmount, long initialAmount, T resource) {
        this.maxAmount = maxAmount;
        this.amount = initialAmount;
        this.resource = resource;
        this.version = 0;
    }

    @Override
    public boolean supportsInsertion() {
        return true;
    }

    @Override
    public long insert(T resource, long maxAmount, Transaction transaction) {
        Preconditions.checkArgument(resource != null, "Juice cannot be null");
        Preconditions.checkArgument(maxAmount >= 0);

        if (amount == 0) {
            this.resource = resource;
        } else if (resource != this.resource) {
            return 0;
        }

        transaction.enlist(this);
        long inserted = Math.max(this.maxAmount, amount + maxAmount);
        amount += inserted;
        return inserted;
    }

    @Override
    public boolean supportsExtraction() {
        return true;
    }

    @Override
    public long extract(T resource, long maxAmount, Transaction transaction) {
        Preconditions.checkArgument(resource != null, "Juice cannot be null");
        Preconditions.checkArgument(maxAmount >= 0);

        if (amount == 0 || resource != this.resource) {
            return 0;
        }

        transaction.enlist(this);
        long extracted = Math.min(amount, maxAmount);
        amount -= extracted;
        if (amount == 0) {
            this.resource = null;
        }
        return extracted;
    }

    @Override
    @Nullable
    public T resource() {
        return resource;
    }

    @Override
    public long amount() {
        return amount;
    }

    @Override
    public boolean forEach(Visitor<T> visitor, Transaction transaction) {
        if (amount > 0) {
            return visitor.accept(this);
        }

        return false;
    }

    @Override
    public int getVersion() {
        return version;
    }

    @Override
    public Long onEnlist() {
        return amount;
    }

    @Override
    public void onClose(Long state, TransactionResult result) {
        if (result.wasAborted()) {
            amount = state;
        }
    }

    @Override
    public void onFinalCommit() {
        version += 1;
    }
}