CommandArguments.java

package redis.clients.jedis;

import java.util.*;

import redis.clients.jedis.annots.Experimental;
import redis.clients.jedis.annots.Internal;
import redis.clients.jedis.args.Rawable;
import redis.clients.jedis.args.RawableFactory;
import redis.clients.jedis.commands.ProtocolCommand;
import redis.clients.jedis.params.IParams;
import redis.clients.jedis.search.RediSearchUtil;
import redis.clients.jedis.util.JedisClusterCRC16;

public class CommandArguments implements Iterable<Rawable> {

  /**
   * Default initial capacity for the keys list. Most Redis commands have 1-3 keys,
   * so a small initial capacity avoids reallocations for common cases.
   */
  private static final int DEFAULT_KEYS_CAPACITY = 4;

  private CommandKeyArgumentPreProcessor keyPreProc = null;
  private final ArrayList<Rawable> args;

  /**
   * Pre-allocated list for storing keys. Using ArrayList directly avoids the
   * memory reallocation overhead of transitioning from emptyList -> singletonList -> ArrayList.
   */
  private final ArrayList<Object> keys;

  /**
   * Cached hash slots computed from keys. Null indicates the cache is invalid
   * and needs to be recomputed. The cache is invalidated when keys are added.
   */
  private Set<Integer> cachedHashSlots;

  private boolean blocking;

  private CommandArguments() {
    throw new InstantiationError();
  }

  public CommandArguments(ProtocolCommand command) {
    args = new ArrayList<>();
    args.add(command);

    keys = new ArrayList<>(DEFAULT_KEYS_CAPACITY);
    cachedHashSlots = null;
  }

  /**
   * Pre-sizes the backing list for the expected total number of arguments, command token
   * included ��� i.e. the RESP array length of the final command, which is also what
   * {@link #size()} returns once fully built. Avoids intermediate {@code ArrayList} growth
   * when commands are built from large collections. Under-estimating is safe: the list
   * falls back to normal growth. Values not above the current capacity have no effect.
   * <p>
   * Example: ZADD with a key and N member/score pairs expects {@code 2 + 2 * N} arguments.
   *
   * @param expectedTotalArguments the expected total argument count, command included
   * @return this
   * @since 8.1
   */
  public CommandArguments ensureCapacity(int expectedTotalArguments) {
    args.ensureCapacity(expectedTotalArguments);
    return this;
  }

  /**
   * Pre-sizes the key-tracking list for the expected total number of keys, mirroring the
   * {@link #ensureCapacity(int)} semantics. Package-private: key counts are only known to
   * the command builders.
   */
  CommandArguments ensureKeyCapacity(int expectedTotalKeys) {
    keys.ensureCapacity(expectedTotalKeys);
    return this;
  }

  public ProtocolCommand getCommand() {
    return (ProtocolCommand) args.get(0);
  }

  @Experimental
  void setKeyArgumentPreProcessor(CommandKeyArgumentPreProcessor keyPreProcessor) {
    this.keyPreProc = keyPreProcessor;
  }

  public CommandArguments add(Rawable arg) {
    args.add(arg);
    return this;
  }

  public CommandArguments add(byte[] arg) {
    return add(RawableFactory.from(arg));
  }

  public CommandArguments add(boolean arg) {
    return add(RawableFactory.from(arg));
  }

  public CommandArguments add(int arg) {
    return add(RawableFactory.from(arg));
  }

  public CommandArguments add(long arg) {
    return add(RawableFactory.from(arg));
  }

  public CommandArguments add(double arg) {
    return add(RawableFactory.from(arg));
  }

  public CommandArguments add(String arg) {
    return add(RawableFactory.from(arg));
  }

  public CommandArguments add(Object arg) {
    if (arg == null) {
      throw new IllegalArgumentException("null is not a valid argument.");
    } else if (arg instanceof Rawable) {
      args.add((Rawable) arg);
    } else if (arg instanceof byte[]) {
      args.add(RawableFactory.from((byte[]) arg));
    } else if (arg instanceof Boolean) {
      args.add(RawableFactory.from((Boolean) arg));
    } else if (arg instanceof Integer) {
      args.add(RawableFactory.from((Integer) arg));
    } else if (arg instanceof Long) {
      args.add(RawableFactory.from((Long) arg));
    } else if (arg instanceof Double) {
      args.add(RawableFactory.from((Double) arg));
    } else if (arg instanceof float[]) {
      args.add(RawableFactory.from(RediSearchUtil.toByteArray((float[]) arg)));
    } else if (arg instanceof String) {
      args.add(RawableFactory.from((String) arg));
    } else if (arg instanceof GeoCoordinate) {
      GeoCoordinate geo = (GeoCoordinate) arg;
      args.add(RawableFactory.from(geo.getLongitude() + "," + geo.getLatitude()));
    } else {
      args.add(RawableFactory.from(String.valueOf(arg)));
    }
    return this;
  }

  public CommandArguments addObjects(Object... args) {
    // Pre-size from the known array length so bulk commands avoid intermediate
    // ArrayList growth copies.
    this.args.ensureCapacity(this.args.size() + args.length);
    for (Object arg : args) {
      add(arg);
    }
    return this;
  }

  public CommandArguments addObjects(Collection args) {
    // Pre-size from the known collection size so bulk commands built from
    // collections (e.g. SADD/ZADD with many members) avoid intermediate
    // ArrayList growth copies.
    this.args.ensureCapacity(this.args.size() + args.size());
    args.forEach(arg -> add(arg));
    return this;
  }

  public CommandArguments key(Object key) {
    if (keyPreProc != null) {
      key = keyPreProc.actualKey(key);
    }

    if (key instanceof Rawable) {
      Rawable raw = (Rawable) key;
      args.add(raw);
      recordKey(raw.getRaw());
    } else if (key instanceof byte[]) {
      Rawable raw = RawableFactory.from((byte[]) key);
      args.add(raw);
      recordKey(raw.getRaw());
    } else if (key instanceof String) {
      String raw = (String) key;
      args.add(RawableFactory.from(raw));
      recordKey(raw);
    } else {
      throw new IllegalArgumentException("\"" + key.toString() + "\" is not a valid argument.");
    }

    return this;
  }

  final CommandArguments addHashSlotKey(String key) {
    return recordKey(key);
  }

  /**
   * Records a routing-only key that is not transmitted. Copied so the recorded key never aliases
   * a caller-owned array.
   */
  final CommandArguments addHashSlotKey(byte[] key) {
    return recordKey(Arrays.copyOf(key, key.length));
  }

  private CommandArguments recordKey(Object key) {
    keys.add(key);
    cachedHashSlots = null;
    return this;
  }

  public final CommandArguments keys(Object... keys) {
    // Pre-size both lists for multi-key commands (DEL, MGET, SINTER, ...) with many keys.
    args.ensureCapacity(args.size() + keys.length);
    ensureKeyCapacity(this.keys.size() + keys.length);
    Arrays.stream(keys).forEach(this::key);
    return this;
  }

  public final CommandArguments keys(Collection keys) {
    args.ensureCapacity(args.size() + keys.size());
    ensureKeyCapacity(this.keys.size() + keys.size());
    keys.forEach(this::key);
    return this;
  }

  public final CommandArguments addParams(IParams params) {
    params.addParams(this);
    return this;
  }

  protected final CommandArguments addHashSlotKeys(byte[]... keys) {
    for (byte[] key : keys) {
      addHashSlotKey(key);
    }
    return this;
  }

  protected final CommandArguments addHashSlotKeys(String... keys) {
    for (String key : keys) {
      addHashSlotKey(key);
    }
    return this;
  }

  public int size() {
    return args.size();
  }

  /**
   * Get the argument at the specified index.
   * @param index the index of the argument to retrieve (0-based, where 0 is the command itself)
   * @return the Rawable argument at the specified index
   * @throws IndexOutOfBoundsException if the index is out of range
   */
  public Rawable get(int index) {
    return args.get(index);
  }

  @Override
  public Iterator<Rawable> iterator() {
    return args.iterator();
  }

  /**
   * Returns the keys used in this command.
   * <p>
   * <b>Internal API:</b> This method is internal and should not be used by external code.
   * It is exposed for internal use by caching ({@link redis.clients.jedis.csc.CacheKey#getRedisKeys()})
   * and cluster operations.
   * <p>
   * <b>Supported types:</b> Keys are stored as either {@link String} or {@code byte[]} depending on
   * how they were added via {@link #key(Object)} or {@link #addHashSlotKey(String)}/{@link #addHashSlotKey(byte[])}.
   * Only {@link String} and {@code byte[]} are guaranteed to be supported by downstream consumers.
   * <p>
   * <b>Type safety:</b> Consumers must handle both {@link String} and {@code byte[]} types.
   * Passing other types may cause {@link IllegalArgumentException} when used with caching
   * (see {@link redis.clients.jedis.csc.AbstractCache#makeKeyForRedisKeysToCacheKeys(Object)})
   * or cluster operations.
   * <p>
   * The returned list is unmodifiable to prevent external modification of the internal key tracking.
   *
   * @return unmodifiable list of keys ({@link String} or {@code byte[]})
   */
  @Internal
  public List<Object> getKeys() {
    return Collections.unmodifiableList(keys);
  }

  @Internal
  public Set<Integer> getKeyHashSlots() {
    // Return cached slots if available (cache is invalidated when keys are added)
    if (cachedHashSlots != null) {
      return cachedHashSlots;
    }

    // Compute hash slots and cache the result
    Set<Integer> slots = new HashSet<>();
    for (Object key : keys) {
      if (key instanceof byte[]) {
        slots.add(JedisClusterCRC16.getSlot((byte[]) key));
      } else {
        slots.add(JedisClusterCRC16.getSlot((String) key));
      }
    }
    // Cache as unmodifiable set to prevent external modification
    cachedHashSlots = Collections.unmodifiableSet(slots);
    return cachedHashSlots;
  }

  /**
   * @return true if this command has no keys, false otherwise
   */
  public boolean isKeyless() {
    return keys.isEmpty();
  }

  public boolean isBlocking() {
    return blocking;
  }

  public CommandArguments blocking() {
    this.blocking = true;
    return this;
  }
}