1 /***********************************************************************
2  * Copyright (c) 2013-2024 Commonwealth Computer Research, Inc.
3  * All rights reserved. This program and the accompanying materials
4  * are made available under the terms of the Apache License, Version 2.0
5  * which accompanies this distribution and is available at
6  * http://www.opensource.org/licenses/apache2.0.php.
7  ***********************************************************************/
8 
9 package org.locationtech.geomesa.utils.io
10 
11 import java.io.InputStream
12 import java.nio.ByteBuffer
13 import java.nio.charset.StandardCharsets
14 
15 object ByteBuffers {
16 
17   class ExpandingByteBuffer(capacity: Int) {
18 
19     private var bb = ByteBuffer.allocate(capacity)
20 
21     private def ensureRemaining(count: Int): Unit = {
22       if (bb.remaining < count) {
23         val expanded = ByteBuffer.allocate(bb.capacity() * 2)
24         bb.flip()
25         expanded.put(bb)
26         bb = expanded
27       }
28     }
29 
30     def putString(string: String): Unit = {
31       if (string == null) { putInt(-1) } else {
32         putBytes(string.getBytes(StandardCharsets.UTF_8))
33       }
34     }
35     def getString: String = {
36       val length = getInt
37       if (length == -1) { null } else {
38         val bytes = Array.ofDim[Byte](length)
39         bb.get(bytes)
40         new String(bytes, StandardCharsets.UTF_8)
41       }
42     }
43 
44     def putBytes(bytes: Array[Byte]): Unit = { ensureRemaining(bytes.length + 4); bb.putBytes(bytes) }
45     def getBytes: Array[Byte] = bb.getBytes
46 
47     def putBool(bool: Boolean): Unit = { ensureRemaining(1); bb.putBool(bool) }
48     def getBool: Boolean = bb.getBool
49 
50     def toArray: Array[Byte] = bb.toArray
51 
52     def put(b: Byte): ByteBuffer = { ensureRemaining(1); bb.put(b) }
53     def get(): Byte = bb.get()
54 
55     def put(src: Array[Byte]): ByteBuffer = { ensureRemaining(src.length); bb.put(src) }
56     def get(dst: Array[Byte]): ByteBuffer = bb.get(dst)
57 
58     def putChar(value: Char): ByteBuffer = { ensureRemaining(2); bb.putChar(value) }
59     def getChar: Char = bb.getChar
60 
61     def putShort(value: Short): ByteBuffer = { ensureRemaining(2); bb.putShort(value) }
62     def getShort: Short = bb.getShort
63 
64     def putInt(value: Int): ByteBuffer = { ensureRemaining(4); bb.putInt(value) }
65     def getInt: Int = bb.getInt
66 
67     def putLong(value: Long): ByteBuffer = { ensureRemaining(8); bb.putLong(value) }
68     def getLong: Long = bb.getLong
69 
70     def putFloat(value: Float): ByteBuffer = { ensureRemaining(4); bb.putFloat(value) }
71     def getFloat: Float = bb.getFloat
72 
73     def putDouble(value: Double): ByteBuffer = { ensureRemaining(8); bb.putDouble(value) }
74     def getDouble: Double = bb.getDouble
75   }
76 
77   implicit class RichByteBuffer(val bb: ByteBuffer) extends AnyVal {
78 
79     def toInputStream: InputStream = new ByteBufferInputStream(bb)
80 
81     def putBytes(bytes: Array[Byte]): Unit = {
82       bb.putInt(bytes.length)
83       bb.put(bytes)
84     }
85     def getBytes: Array[Byte] = {
86       val bytes = Array.ofDim[Byte](bb.getInt())
87       bb.get(bytes)
88       bytes
89     }
90 
91     def putString(string: String): Unit = {
92       if (string == null) { bb.putInt(-1) } else {
93         putBytes(string.getBytes(StandardCharsets.UTF_8))
94       }
95     }
96     def getString: String = {
97       val length = bb.getInt
98       if (length == -1) { null } else {
99         val bytes = Array.ofDim[Byte](length)
100         bb.get(bytes)
101         new String(bytes, StandardCharsets.UTF_8)
102       }
103     }
104 
105     def putBool(bool: Boolean): Unit = bb.put(if (bool) { 1.toByte } else { 0.toByte })
106     def getBool: Boolean = bb.get == 1
107 
108     def toArray: Array[Byte] = {
109       bb.flip()
110       val bytes = Array.ofDim[Byte](bb.remaining)
111       bb.get(bytes)
112       bytes
113     }
114   }
115 
116   class ByteBufferInputStream(buffer: ByteBuffer) extends InputStream {
117 
118     override def read(): Int = {
119       if (!buffer.hasRemaining) { -1 } else {
120         buffer.get() & 0xFF
121       }
122     }
123 
124     override def read(bytes: Array[Byte], offset: Int, length: Int): Int = {
125       if (!buffer.hasRemaining) { -1 } else {
126         val read = math.min(length, buffer.remaining)
127         buffer.get(bytes, offset, read)
128         read
129       }
130     }
131   }
132 
133 }