<commit_before>/*
 * 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.jackrabbit.oak.plugins.segment;

import static com.google.common.base.Objects.equal;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkPositionIndexes;
import static org.apache.jackrabbit.oak.plugins.segment.SegmentWriter.BLOCK_SIZE;

import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.Callable;

import org.apache.jackrabbit.oak.api.PropertyState;
import org.apache.jackrabbit.oak.api.Type;
import org.apache.jackrabbit.oak.plugins.memory.PropertyStates;

import com.google.common.base.Charsets;
import com.google.common.cache.Weigher;

public class Segment {

    /**
     * Number of bytes used for storing a record identifier. One byte
     * is used for identifying the segment and two for the record offset
     * within that segment.
     */
    static final int RECORD_ID_BYTES = 1 + 2;

    /**
     * The limit on segment references within one segment. Since record
     * identifiers use one byte to indicate the referenced segment, a single
     * segment can hold references to up to 256 segments.
     */
    static final int SEGMENT_REFERENCE_LIMIT = 1 << 8; // 256

    /**
     * The number of bytes (or bits of address space) to use for the
     * alignment boundary of segment records.
     */
    static final int RECORD_ALIGN_BITS = 2;
    static final int RECORD_ALIGN_BYTES = 1 << RECORD_ALIGN_BITS; // 4

    /**
     * Maximum segment size. Record identifiers are stored as three-byte
     * sequences with the first byte indicating the segment and the next
     * two the offset within that segment. Since all records are aligned
     * at four-byte boundaries, the two bytes can address up to 256kB of
     * record data.
     */
    static final int MAX_SEGMENT_SIZE = 1 << (16 + RECORD_ALIGN_BITS); // 256kB

    /**
     * The size limit for small values. The variable length of small values
     * is encoded as a single byte with the high bit as zero, which gives us
     * seven bits for encoding the length of the value.
     */
    static final int SMALL_LIMIT = 1 << 7;

    /**
     * The size limit for medium values. The variable length of medium values
     * is encoded as two bytes with the highest bits of the first byte set to
     * one and zero, which gives us 14 bits for encoding the length of the
     * value. And since small values are never stored as medium ones, we can
     * extend the size range to cover that many longer values.
     */
    static final int MEDIUM_LIMIT = (1 << (16 - 2)) + SMALL_LIMIT;

    static final Weigher<UUID, Segment> WEIGHER =
            new Weigher<UUID, Segment>() {
                @Override
                public int weigh(UUID key, Segment value) {
                    return value.size();
                }
            };

    final SegmentStore store; // TODO: should be private

    private final UUID uuid;

    private final ByteBuffer data;

    private final List<UUID> uuids;

    public Segment(
            SegmentStore store, UUID uuid, ByteBuffer data, List<UUID> uuids) {
        this.store = checkNotNull(store);
        this.uuid = checkNotNull(uuid);
        this.data = checkNotNull(data);
        this.uuids = checkNotNull(uuids);
    }

    /**
     * Maps the given record offset to the respective position within the
     * internal {@link #data} array. The validity of a record with the given
     * length at the given offset is also verified.
     *
     * @param offset record offset
     * @param length record length
     * @return position within the data array
     */
    private int pos(int offset, int length) {
        int pos = offset - (MAX_SEGMENT_SIZE - size());
        checkPositionIndexes(pos, pos + length, size());
        return data.position() + pos;
    }

    public UUID getSegmentId() {
        return uuid;
    }

    public ByteBuffer getData() {
        return data;
    }

    public int size() {
        return data.remaining();
    }

    byte readByte(int offset) {
        return data.get(pos(offset, 1));
    }

    /**
     * Returns the identified segment.
     *
     * @param uuid segment identifier
     * @return identified segment
     */
    Segment getSegment(UUID uuid) {
        if (equal(uuid, this.uuid)) {
            return this; // optimization for the common case (OAK-1031)
        } else {
            return store.readSegment(uuid);
        }
    }

    /**
     * Returns the segment that contains the identified record.
     *
     * @param id record identifier
     * @return segment that contains the identified record
     */
    Segment getSegment(RecordId id) {
        return getSegment(checkNotNull(id).getSegmentId());
    }


    /**
     * Reads the given number of bytes starting from the given position
     * in this segment.
     *
     * @param position position within segment
     * @param buffer target buffer
     * @param offset offset within target buffer
     * @param length number of bytes to read
     */
     void readBytes(int position, byte[] buffer, int offset, int length) {
        checkNotNull(buffer);
        checkPositionIndexes(offset, offset + length, buffer.length);
        ByteBuffer d = data.duplicate();
        d.position(pos(position, length));
        d.get(buffer, offset, length);
    }

    RecordId readRecordId(int offset) {
        int pos = pos(offset, RECORD_ID_BYTES);
        return internalReadRecordId(pos);
    }

    private RecordId internalReadRecordId(int pos) {
        return new RecordId(
                uuids.get(data.get(pos) & 0xff),
                (data.get(pos + 1) & 0xff) << (8 + Segment.RECORD_ALIGN_BITS)
                | (data.get(pos + 2) & 0xff) << Segment.RECORD_ALIGN_BITS);
    }

    int readInt(int offset) {
        int pos = pos(offset, 4);
        return (data.get(pos) & 0xff) << 24
                | (data.get(pos + 1) & 0xff) << 16
                | (data.get(pos + 2) & 0xff) << 8
                | (data.get(pos + 3) & 0xff);
    }

    String readString(final RecordId id) {
        return store.getRecord(id, new Callable<String>() {
            @Override
            public String call() {
                return getSegment(id).readString(id.getOffset());
            }
        });
    }

    String readString(int offset) {
        int pos = pos(offset, 1);
        long length = internalReadLength(pos);
        if (length < SMALL_LIMIT) {
            byte[] bytes = new byte[(int) length];
            ByteBuffer buffer = data.duplicate();
            buffer.position(pos + 1);
            buffer.get(bytes);
            return new String(bytes, Charsets.UTF_8);
        } else if (length < MEDIUM_LIMIT) {
            byte[] bytes = new byte[(int) length];
            ByteBuffer buffer = data.duplicate();
            buffer.position(pos + 2);
            buffer.get(bytes);
            return new String(bytes, Charsets.UTF_8);
        } else if (length < Integer.MAX_VALUE) {
            int size = (int) ((length + BLOCK_SIZE - 1) / BLOCK_SIZE);
            ListRecord list =
                    new ListRecord(this, internalReadRecordId(pos + 8), size);
            SegmentStream stream = new SegmentStream(
                    store, new RecordId(uuid, offset), list, length);
            try {
                return stream.getString();
            } finally {
                stream.close();
            }
        } else {
            throw new IllegalStateException("String is too long: " + length);
        }
    }

    MapRecord readMap(final RecordId id) {
        return store.getRecord(id, new Callable<MapRecord>() {
            @Override
            public MapRecord call() {
                return getSegment(id).readMap(id.getOffset());
            }
        });
    }

    MapRecord readMap(int offset) {
        int head = readInt(offset);
        int level = head >>> MapRecord.SIZE_BITS;
        int size = head & ((1 << MapRecord.SIZE_BITS) - 1);
        if (size > MapRecord.BUCKETS_PER_LEVEL
                && level < MapRecord.MAX_NUMBER_OF_LEVELS) {
            int bitmap = readInt(offset + 4);
            return new MapBranch(this, offset, size, level, bitmap);
        } else {
            return new MapLeaf(this, offset, size, level);
        }
    }

    Template readTemplate(final RecordId id) {
        return store.getRecord(id, new Callable<Template>() {
            @Override
            public Template call() {
                return getSegment(id).readTemplate(id.getOffset());
            }
        });
    }

    Template readTemplate(int offset) {
        int head = readInt(offset);
        boolean hasPrimaryType = (head & (1 << 31)) != 0;
        boolean hasMixinTypes = (head & (1 << 30)) != 0;
        boolean zeroChildNodes = (head & (1 << 29)) != 0;
        boolean manyChildNodes = (head & (1 << 28)) != 0;<commit_msg>OAK-1031: SegmentMK: Fewer segment lookups<commit_after>/*
 * 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.jackrabbit.oak.plugins.segment;

import static com.google.common.base.Objects.equal;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkPositionIndexes;
import static org.apache.jackrabbit.oak.plugins.segment.SegmentWriter.BLOCK_SIZE;

import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.Callable;

import org.apache.jackrabbit.oak.api.PropertyState;
import org.apache.jackrabbit.oak.api.Type;
import org.apache.jackrabbit.oak.plugins.memory.PropertyStates;

import com.google.common.base.Charsets;
import com.google.common.cache.Weigher;

public class Segment {

    /**
     * Number of bytes used for storing a record identifier. One byte
     * is used for identifying the segment and two for the record offset
     * within that segment.
     */
    static final int RECORD_ID_BYTES = 1 + 2;

    /**
     * The limit on segment references within one segment. Since record
     * identifiers use one byte to indicate the referenced segment, a single
     * segment can hold references to up to 256 segments.
     */
    static final int SEGMENT_REFERENCE_LIMIT = 1 << 8; // 256

    /**
     * The number of bytes (or bits of address space) to use for the
     * alignment boundary of segment records.
     */
    static final int RECORD_ALIGN_BITS = 2;
    static final int RECORD_ALIGN_BYTES = 1 << RECORD_ALIGN_BITS; // 4

    /**
     * Maximum segment size. Record identifiers are stored as three-byte
     * sequences with the first byte indicating the segment and the next
     * two the offset within that segment. Since all records are aligned
     * at four-byte boundaries, the two bytes can address up to 256kB of
     * record data.
     */
    static final int MAX_SEGMENT_SIZE = 1 << (16 + RECORD_ALIGN_BITS); // 256kB

    /**
     * The size limit for small values. The variable length of small values
     * is encoded as a single byte with the high bit as zero, which gives us
     * seven bits for encoding the length of the value.
     */
    static final int SMALL_LIMIT = 1 << 7;

    /**
     * The size limit for medium values. The variable length of medium values
     * is encoded as two bytes with the highest bits of the first byte set to
     * one and zero, which gives us 14 bits for encoding the length of the
     * value. And since small values are never stored as medium ones, we can
     * extend the size range to cover that many longer values.
     */
    static final int MEDIUM_LIMIT = (1 << (16 - 2)) + SMALL_LIMIT;

    static final Weigher<UUID, Segment> WEIGHER =
            new Weigher<UUID, Segment>() {
                @Override
                public int weigh(UUID key, Segment value) {
                    return value.size();
                }
            };

    final SegmentStore store; // TODO: should be private

    private final UUID uuid;

    private final ByteBuffer data;

    private final List<UUID> uuids;

    public Segment(
            SegmentStore store, UUID uuid, ByteBuffer data, List<UUID> uuids) {
        this.store = checkNotNull(store);
        this.uuid = checkNotNull(uuid);
        this.data = checkNotNull(data);
        this.uuids = checkNotNull(uuids);
    }

    /**
     * Maps the given record offset to the respective position within the
     * internal {@link #data} array. The validity of a record with the given
     * length at the given offset is also verified.
     *
     * @param offset record offset
     * @param length record length
     * @return position within the data array
     */
    private int pos(int offset, int length) {
        int pos = offset - (MAX_SEGMENT_SIZE - size());
        checkPositionIndexes(pos, pos + length, size());
        return data.position() + pos;
    }

    public UUID getSegmentId() {
        return uuid;
    }

    public ByteBuffer getData() {
        return data;
    }

    public int size() {
        return data.remaining();
    }

    byte readByte(int offset) {
        return data.get(pos(offset, 1));
    }

    /**
     * Returns the identified segment.
     *
     * @param uuid segment identifier
     * @return identified segment
     */
    Segment getSegment(UUID uuid) {
        if (equal(uuid, this.uuid)) {
            return this; // optimization for the common case (OAK-1031)
        } else {
            return store.readSegment(uuid);
        }
    }

    /**
     * Returns the segment that contains the identified record.
     *
     * @param id record identifier
     * @return segment that contains the identified record
     */
    Segment getSegment(RecordId id) {
        return getSegment(checkNotNull(id).getSegmentId());
    }


    /**
     * Reads the given number of bytes starting from the given position
     * in this segment.
     *
     * @param position position within segment
     * @param buffer target buffer
     * @param offset offset within target buffer
     * @param length number of bytes to read
     */
     void readBytes(int position, byte[] buffer, int offset, int length) {
        checkNotNull(buffer);
        checkPositionIndexes(offset, offset + length, buffer.length);
        ByteBuffer d = data.duplicate();
        d.position(pos(position, length));
        d.get(buffer, offset, length);
    }

    RecordId readRecordId(int offset) {
        int pos = pos(offset, RECORD_ID_BYTES);
        return internalReadRecordId(pos);
    }

    private RecordId internalReadRecordId(int pos) {
        return new RecordId(
                uuids.get(data.get(pos) & 0xff),
                (data.get(pos + 1) & 0xff) << (8 + Segment.RECORD_ALIGN_BITS)
                | (data.get(pos + 2) & 0xff) << Segment.RECORD_ALIGN_BITS);
    }

    int readInt(int offset) {
        int pos = pos(offset, 4);
        return (data.get(pos) & 0xff) << 24
                | (data.get(pos + 1) & 0xff) << 16
                | (data.get(pos + 2) & 0xff) << 8
                | (data.get(pos + 3) & 0xff);
    }

    String readString(final RecordId id) {
        return store.getRecord(id, new Callable<String>() {
            @Override
            public String call() {
                return getSegment(id).readString(id.getOffset());
            }
        });
    }

    String readString(int offset) {
        int pos = pos(offset, 1);
        long length = internalReadLength(pos);
        if (length < SMALL_LIMIT) {
            byte[] bytes = new byte[(int) length];
            ByteBuffer buffer = data.duplicate();
            buffer.position(pos + 1);
            buffer.get(bytes);
            return new String(bytes, Charsets.UTF_8);
        } else if (length < MEDIUM_LIMIT) {
            byte[] bytes = new byte[(int) length];
            ByteBuffer buffer = data.duplicate();
            buffer.position(pos + 2);
            buffer.get(bytes);
            return new String(bytes, Charsets.UTF_8);
        } else if (length < Integer.MAX_VALUE) {
            int size = (int) ((length + BLOCK_SIZE - 1) / BLOCK_SIZE);
            ListRecord list =
                    new ListRecord(this, internalReadRecordId(pos + 8), size);
            SegmentStream stream = new SegmentStream(
                    store, new RecordId(uuid, offset), list, length);
            try {
                return stream.getString();
            } finally {
                stream.close();
            }
        } else {
            throw new IllegalStateException("String is too long: " + length);
        }
    }

    MapRecord readMap(RecordId id) {
        return getSegment(id).readMap(id.getOffset());
    }

    MapRecord readMap(int offset) {
        int head = readInt(offset);
        int level = head >>> MapRecord.SIZE_BITS;
        int size = head & ((1 << MapRecord.SIZE_BITS) - 1);
        if (size > MapRecord.BUCKETS_PER_LEVEL
                && level < MapRecord.MAX_NUMBER_OF_LEVELS) {
            int bitmap = readInt(offset + 4);
            return new MapBranch(this, offset, size, level, bitmap);
        } else {
            return new MapLeaf(this, offset, size, level);
        }
    }

    Template readTemplate(final RecordId id) {
        return store.getRecord(id, new Callable<Template>() {
            @Override
            public Template call() {
                return getSegment(id).readTemplate(id.getOffset());
            }
        });
    }

    Template readTemplate(int offset) {
        int head = readInt(offset);
        boolean hasPrimaryType = (head & (1 << 31)) != 0;
        boolean hasMixinTypes = (head & (1 << 30)) != 0;
        boolean zeroChildNodes = (head & (1 << 29)) != 0;
        boolean manyChildNodes = (head & (1 << 28)) != 0;
{'fraction_non_alphanumeric': 0.09393437532346548, 'fraction_numerical': 0.00890177000310527, 'mean_word_length': 2.255222371967655, 'pattern_counts': {'":': 0, '<': 84, '<?xml version=': 0, '>': 24, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6442793', 'n_tokens_mistral': 5459, 'n_tokens_neox': 5214, 'n_words': 1818}