/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache license, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the license for the specific language governing permissions and
 * limitations under the license.
 */
package org.apache.logging.log4j.util;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InvalidObjectException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.StreamCorruptedException;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
import java.util.Arrays;
import java.util.Collection;
import java.util.ConcurrentModificationException;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;

import org.apache.logging.log4j.status.StatusLogger;

Consider this class private. Array-based implementation of the ReadOnlyStringMap interface. Keys are held in a sorted array. 

This is not a generic collection, but makes some trade-offs to optimize for the Log4j context data use case:


  
Garbage-free iteration over key-value pairs with BiConsumer and TriConsumer.
  
Fast copy. If the ThreadContextMap is also an instance of SortedArrayStringMap, the full thread context data can be transferred with two array copies and two field updates.
  
Acceptable performance for small data sets. The current implementation stores keys in a sorted array, values are stored in a separate array at the same index. Worst-case performance of get and containsKey is O(log N), worst-case performance of put and remove is O(N log N). The expectation is that for the small values of N (less than 100) that are the vast majority of ThreadContext use cases, the constants dominate performance more than the asymptotic performance of the algorithms used. 
    
Compact representation.

Since:
2.7/**
 * <em>Consider this class private.</em>
 * Array-based implementation of the {@code ReadOnlyStringMap} interface. Keys are held in a sorted array.
 * <p>
 * This is not a generic collection, but makes some trade-offs to optimize for the Log4j context data use case:
 * </p>
 * <ul>
 *   <li>Garbage-free iteration over key-value pairs with {@code BiConsumer} and {@code TriConsumer}.</li>
 *   <li>Fast copy. If the ThreadContextMap is also an instance of {@code SortedArrayStringMap}, the full thread context
 *     data can be transferred with two array copies and two field updates.</li>
 *   <li>Acceptable performance for small data sets. The current implementation stores keys in a sorted array, values
 *     are stored in a separate array at the same index.
 *     Worst-case performance of {@code get} and {@code containsKey} is O(log N),
 *     worst-case performance of {@code put} and {@code remove} is O(N log N).
 *     The expectation is that for the small values of {@code N} (less than 100) that are the vast majority of
 *     ThreadContext use cases, the constants dominate performance more than the asymptotic performance of the
 *     algorithms used.
 *     </li>
 *     <li>Compact representation.</li>
 * </ul>
 *
 * @since 2.7
 */
public class SortedArrayStringMap implements IndexedStringMap {

    
The default initial capacity.
/**
     * The default initial capacity.
     */
    private static final int DEFAULT_INITIAL_CAPACITY = 4;
    private static final long serialVersionUID = -5748905872274478116L;
    private static final int HASHVAL = 31;

    private static final TriConsumer<String, Object, StringMap> PUT_ALL = new TriConsumer<String, Object, StringMap>() {
        @Override
        public void accept(final String key, final Object value, final StringMap contextData) {
            contextData.putValue(key, value);
        }
    };

    
An empty array instance to share when the table is not inflated.
/**
     * An empty array instance to share when the table is not inflated.
     */
    private static final String[] EMPTY = {};
    private static final String FROZEN = "Frozen collection cannot be modified";

    private transient String[] keys = EMPTY;
    private transient Object[] values = EMPTY;

    
The number of key-value mappings contained in this map.
/**
     * The number of key-value mappings contained in this map.
     */
    private transient int size;

    private static final Method setObjectInputFilter;
    private static final Method getObjectInputFilter;
    private static final Method newObjectInputFilter;

    static {
        Method[] methods = ObjectInputStream.class.getMethods();
        Method setMethod = null;
        Method getMethod = null;
        for (final Method method : methods) {
            if (method.getName().equals("setObjectInputFilter")) {
                setMethod = method;
            } else if (method.getName().equals("getObjectInputFilter")) {
                getMethod = method;
            }
        }
        Method newMethod = null;
        try {
            if (setMethod != null) {
                final Class<?> clazz = Class.forName("org.apache.logging.log4j.util.internal.DefaultObjectInputFilter");
                methods = clazz.getMethods();
                for (final Method method : methods) {
                    if (method.getName().equals("newInstance") && Modifier.isStatic(method.getModifiers())) {
                        newMethod = method;
                        break;
                    }
                }
            }
        } catch (final ClassNotFoundException ex) {
            // Ignore the exception
        }
        newObjectInputFilter = newMethod;
        setObjectInputFilter = setMethod;
        getObjectInputFilter = getMethod;
    }

    
The next size value at which to resize (capacity * load factor).
@serial
/**
     * The next size value at which to resize (capacity * load factor).
     * @serial
     */
    // If table == EMPTY_TABLE then this is the initial capacity at which the
    // table will be created when inflated.
    private int threshold;
    private boolean immutable;
    private transient boolean iterating;

    public SortedArrayStringMap() {
        this(DEFAULT_INITIAL_CAPACITY);
    }

    public SortedArrayStringMap(final int initialCapacity) {
        if (initialCapacity < 0) {
            throw new IllegalArgumentException("Initial capacity must be at least zero but was " + initialCapacity);
        }
        threshold = ceilingNextPowerOfTwo(initialCapacity == 0 ? 1 : initialCapacity);
    }

    public SortedArrayStringMap(final ReadOnlyStringMap other) {
        if (other instanceof SortedArrayStringMap) {
            initFrom0((SortedArrayStringMap) other);
        } else if (other != null) {
            resize(ceilingNextPowerOfTwo(other.size()));
            other.forEach(PUT_ALL, this);
        }
    }

    public SortedArrayStringMap(final Map<String, ?> map) {
        resize(ceilingNextPowerOfTwo(map.size()));
        for (final Map.Entry<String, ?> entry : map.entrySet()) {
            putValue(entry.getKey(), entry.getValue());
        }
    }

    private void assertNotFrozen() {
        if (immutable) {
            throw new UnsupportedOperationException(FROZEN);
        }
    }

    private void assertNoConcurrentModification() {
        if (iterating) {
            throw new ConcurrentModificationException();
        }
    }

    @Override
    public void clear() {
        if (keys == EMPTY) {
            return;
        }
        assertNotFrozen();
        assertNoConcurrentModification();

        Arrays.fill(keys, 0, size, null);
        Arrays.fill(values, 0, size, null);
        size = 0;
    }

    @Override
    public boolean containsKey(final String key) {
        return indexOfKey(key) >= 0;
    }

    @Override
    public Map<String, String> toMap() {
        final Map<String, String> result = new HashMap<>(size());
        for (int i = 0; i < size(); i++) {
            final Object value = getValueAt(i);
            result.put(getKeyAt(i), value == null ? null : String.valueOf(value));
        }
        return result;
    }

    @Override
    public void freeze() {
        immutable = true;
    }

    @Override
    public boolean isFrozen() {
        return immutable;
    }

    @SuppressWarnings("unchecked")
    @Override
    public <V> V getValue(final String key) {
        final int index = indexOfKey(key);
        if (index < 0) {
            return null;
        }
        return (V) values[index];
    }

    @Override
    public boolean isEmpty() {
        return size == 0;
    }

    @Override
    public int indexOfKey(final String key) {
        if (keys == EMPTY) {
            return -1;
        }
        if (key == null) { // null key is located at the start of the array
            return nullKeyIndex(); // insert at index zero
        }
        final int start = size > 0 && keys[0] == null ? 1 : 0;
        return Arrays.binarySearch(keys, start, size, key);
    }

    private int nullKeyIndex() {
        return size > 0 && keys[0] == null ? 0 : ~0;
    }

    @Override
    public void putValue(final String key, final Object value) {
        assertNotFrozen();
        assertNoConcurrentModification();

        if (keys == EMPTY) {
            inflateTable(threshold);
        }
        final int index = indexOfKey(key);
        if (index >= 0) {
            keys[index] = key;
            values[index] = value;
        } else { // not found, so insert.
            insertAt(~index, key, value);
        }
    }

    private void insertAt(final int index, final String key, final Object value) {
        ensureCapacity();
        System.arraycopy(keys, index, keys, index + 1, size - index);
        System.arraycopy(values, index, values, index + 1, size - index);
        keys[index] = key;
        values[index] = value;
        size++;
    }

    @Override
    public void putAll(final ReadOnlyStringMap source) {
        if (source == this || source == null || source.isEmpty()) {
            // this.putAll(this) does not modify this collection
            // this.putAll(null) does not modify this collection
            // this.putAll(empty ReadOnlyStringMap) does not modify this collection
            return;
        }
        assertNotFrozen();
        assertNoConcurrentModification();

        if (source instanceof SortedArrayStringMap) {
            if (this.size == 0) {
                initFrom0((SortedArrayStringMap) source);
            } else {
                merge((SortedArrayStringMap) source);
            }
        } else if (source != null) {
            source.forEach(PUT_ALL, this);
        }
    }

    private void initFrom0(final SortedArrayStringMap other) {
        if (keys.length < other.size) {
            keys = new String[other.threshold];
            values = new Object[other.threshold];
        }
        System.arraycopy(other.keys, 0, keys, 0, other.size);
        System.arraycopy(other.values, 0, values, 0, other.size);

        size = other.size;
        threshold = other.threshold;
    }

    private void merge(final SortedArrayStringMap other) {
        final String[] myKeys = keys;
        final Object[] myVals = values;
        final int newSize = other.size + this.size;
        threshold = ceilingNextPowerOfTwo(newSize);
        if (keys.length < threshold) {
            keys = new String[threshold];
            values = new Object[threshold];
        }
        // move largest collection to the beginning of this data structure, smallest to the end
        boolean overwrite = true;
        if (other.size() > size()) {
            // move original key-values to end
            System.arraycopy(myKeys, 0, keys, other.size, this.size);
            System.arraycopy(myVals, 0, values, other.size, this.size);

            // insert additional key-value pairs at the beginning
            System.arraycopy(other.keys, 0, keys, 0, other.size);
            System.arraycopy(other.values, 0, values, 0, other.size);
            size = other.size;

            // loop over original keys and insert (drop values for same key)
            overwrite = false;
        } else {
            System.arraycopy(myKeys, 0, keys, 0, this.size);
            System.arraycopy(myVals, 0, values, 0, this.size);

            // move additional key-value pairs to end
            System.arraycopy(other.keys, 0, keys, this.size, other.size);
            System.arraycopy(other.values, 0, values, this.size, other.size);

            // new values are at the end, will be processed below. Overwrite is true.
        }
        for (int i = size; i < newSize; i++) {
            final int index = indexOfKey(keys[i]);
            if (index < 0) { // key does not exist
                insertAt(~index, keys[i], values[i]);
            } else if (overwrite) { // existing key: only overwrite when looping over the new values
                keys[index] = keys[i];
                values[index] = values[i];
            }
        }
        // prevent memory leak: null out references
        Arrays.fill(keys, size, newSize, null);
        Arrays.fill(values, size, newSize, null);
    }

    private void ensureCapacity() {
        if (size >= threshold) {
            resize(threshold * 2);
        }
    }

    private void resize(final int newCapacity) {
        final String[] oldKeys = keys;
        final Object[] oldValues = values;

        keys = new String[newCapacity];
        values = new Object[newCapacity];

        System.arraycopy(oldKeys, 0, keys, 0, size);
        System.arraycopy(oldValues, 0, values, 0, size);

        threshold = newCapacity;
    }

    
Inflates the table.
/**
     * Inflates the table.
     */
    private void inflateTable(final int toSize) {
        threshold = toSize;
        keys = new String[toSize];
        values = new Object[toSize];
    }

    @Override
    public void remove(final String key) {
        if (keys == EMPTY) {
            return;
        }

        final int index = indexOfKey(key);
        if (index >= 0) {
            assertNotFrozen();
            assertNoConcurrentModification();

            System.arraycopy(keys, index + 1, keys, index, size - 1 - index);
            System.arraycopy(values, index + 1, values, index, size - 1 - index);
            keys[size - 1] = null;
            values[size - 1] = null;
            size--;
        }
    }

    @Override
    public String getKeyAt(final int index) {
        if (index < 0 || index >= size) {
            return null;
        }
        return keys[index];
    }

    @SuppressWarnings("unchecked")
    @Override
    public <V> V getValueAt(final int index) {
        if (index < 0 || index >= size) {
            return null;
        }
        return (V) values[index];
    }

    @Override
    public int size() {
        return size;
    }

    @SuppressWarnings("unchecked")
    @Override
    public <V> void forEach(final BiConsumer<String, ? super V> action) {
        iterating = true;
        try {
            for (int i = 0; i < size; i++) {
                action.accept(keys[i], (V) values[i]);
            }
        } finally {
            iterating = false;
        }
    }

    @SuppressWarnings("unchecked")
    @Override
    public <V, T> void forEach(final TriConsumer<String, ? super V, T> action, final T state) {
        iterating = true;
        try {
            for (int i = 0; i < size; i++) {
                action.accept(keys[i], (V) values[i], state);
            }
        } finally {
            iterating = false;
        }
    }

    @Override
    public boolean equals(final Object obj) {
        if (obj == this) {
            return true;
        }
        if (!(obj instanceof SortedArrayStringMap)) {
            return false;
        }
        final SortedArrayStringMap other = (SortedArrayStringMap) obj;
        if (this.size() != other.size()) {
            return false;
        }
        for (int i = 0; i < size(); i++) {
            if (!Objects.equals(keys[i], other.keys[i])) {
                return false;
            }
            if (!Objects.equals(values[i], other.values[i])) {
                return false;
            }
        }
        return true;
    }

    @Override
    public int hashCode() {
        int result = 37;
        result = HASHVAL * result + size;
        result = HASHVAL * result + hashCode(keys, size);
        result = HASHVAL * result + hashCode(values, size);
        return result;
    }

    private static int hashCode(final Object[] values, final int length) {
        int result = 1;
        for (int i = 0; i < length; i++) {
            result = HASHVAL * result + (values[i] == null ? 0 : values[i].hashCode());
        }
        return result;
    }

    @Override
    public String toString() {
        final StringBuilder sb = new StringBuilder(256);
        sb.append('{');
        for (int i = 0; i < size; i++) {
            if (i > 0) {
                sb.append(", ");
            }
            sb.append(keys[i]).append('=');
            sb.append(values[i] == this ? "(this map)" : values[i]);
        }
        sb.append('}');
        return sb.toString();
    }

    
Save the state of the SortedArrayStringMap instance to a stream (i.e., serialize it). 
@serialData
The capacity of the SortedArrayStringMap (the length of the
            bucket array) is emitted (int), followed by the
            size (an int, the number of key-value
            mappings), followed by the key (Object) and value (Object)
            for each key-value mapping.  The key-value mappings are
            emitted in no particular order./**
     * Save the state of the {@code SortedArrayStringMap} instance to a stream (i.e.,
     * serialize it).
     *
     * @serialData The <i>capacity</i> of the SortedArrayStringMap (the length of the
     *             bucket array) is emitted (int), followed by the
     *             <i>size</i> (an int, the number of key-value
     *             mappings), followed by the key (Object) and value (Object)
     *             for each key-value mapping.  The key-value mappings are
     *             emitted in no particular order.
     */
    private void writeObject(final java.io.ObjectOutputStream s) throws IOException {
        // Write out the threshold, and any hidden stuff
        s.defaultWriteObject();

        // Write out number of buckets
        if (keys == EMPTY) {
            s.writeInt(ceilingNextPowerOfTwo(threshold));
        } else {
            s.writeInt(keys.length);
        }

        // Write out size (number of Mappings)
        s.writeInt(size);

        // Write out keys and values (alternating)
        if (size > 0) {
            for (int i = 0; i < size; i++) {
                s.writeObject(keys[i]);
                try {
                    s.writeObject(marshall(values[i]));
                } catch (final Exception e) {
                    handleSerializationException(e, i, keys[i]);
                    s.writeObject(null);
                }
            }
        }
    }

    private static byte[] marshall(final Object obj) throws IOException {
        if (obj == null) {
            return null;
        }
        final ByteArrayOutputStream bout = new ByteArrayOutputStream();
        try (ObjectOutputStream oos = new ObjectOutputStream(bout)) {
            oos.writeObject(obj);
            oos.flush();
            return bout.toByteArray();
        }
    }

    private static Object unmarshall(final byte[] data, final ObjectInputStream inputStream)
            throws IOException, ClassNotFoundException {
        final ByteArrayInputStream bin = new ByteArrayInputStream(data);
        Collection<String> allowedClasses = null;
        ObjectInputStream ois;
        if (inputStream instanceof FilteredObjectInputStream) {
            allowedClasses = ((FilteredObjectInputStream) inputStream).getAllowedClasses();
            ois = new FilteredObjectInputStream(bin, allowedClasses);
        } else {
            try {
                final Object obj = getObjectInputFilter.invoke(inputStream);
                final Object filter = newObjectInputFilter.invoke(null, obj);
                ois = new ObjectInputStream(bin);
                setObjectInputFilter.invoke(ois, filter);
            } catch (IllegalAccessException | InvocationTargetException ex) {
                throw new StreamCorruptedException("Unable to set ObjectInputFilter on stream");
            }
        }
        try {
            return ois.readObject();
        } finally {
            ois.close();
        }
    }

    
Calculate the next power of 2, greater than or equal to x.

From Hacker's Delight, Chapter 3, Harry S. Warren Jr.
Params:
x – Value to round up
Returns:
The next power of 2 from x inclusive/**
     * Calculate the next power of 2, greater than or equal to x.
     * <p>
     * From Hacker's Delight, Chapter 3, Harry S. Warren Jr.
     *
     * @param x Value to round up
     * @return The next power of 2 from x inclusive
     */
    private static int ceilingNextPowerOfTwo(final int x) {
        final int BITS_PER_INT = 32;
        return 1 << (BITS_PER_INT - Integer.numberOfLeadingZeros(x - 1));
    }

    
Reconstitute the SortedArrayStringMap instance from a stream (i.e., deserialize it). /**
     * Reconstitute the {@code SortedArrayStringMap} instance from a stream (i.e.,
     * deserialize it).
     */
    private void readObject(final java.io.ObjectInputStream s)  throws IOException, ClassNotFoundException {
        if (!(s instanceof FilteredObjectInputStream) && setObjectInputFilter == null) {
            throw new IllegalArgumentException("readObject requires a FilteredObjectInputStream or an ObjectInputStream that accepts an ObjectInputFilter");
        }
        // Read in the threshold (ignored), and any hidden stuff
        s.defaultReadObject();

        // set other fields that need values
        keys = EMPTY;
        values = EMPTY;

        // Read in number of buckets
        final int capacity = s.readInt();
        if (capacity < 0) {
            throw new InvalidObjectException("Illegal capacity: " + capacity);
        }

        // Read number of mappings
        final int mappings = s.readInt();
        if (mappings < 0) {
            throw new InvalidObjectException("Illegal mappings count: " + mappings);
        }

        // allocate the bucket array;
        if (mappings > 0) {
            inflateTable(capacity);
        } else {
            threshold = capacity;
        }

        // Read the keys and values, and put the mappings in the arrays
        for (int i = 0; i < mappings; i++) {
            keys[i] = (String) s.readObject();
            try {
                final byte[] marshalledObject = (byte[]) s.readObject();
                values[i] = marshalledObject == null ? null : unmarshall(marshalledObject, s);
            } catch (final Exception | LinkageError error) {
                handleSerializationException(error, i, keys[i]);
                values[i] = null;
            }
        }
        size = mappings;
    }

    private void handleSerializationException(final Throwable t, final int i, final String key) {
        StatusLogger.getLogger().warn("Ignoring {} for key[{}] ('{}')", String.valueOf(t), i, keys[i]);
    }
}