buffer.rs 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337
  1. use lua_shared as lua;
  2. use lua_shared::lua_State;
  3. use paste::paste;
  4. use std::u64;
  5. use crate::{check_slice, insert_function};
  6. macro_rules! def_rw {
  7. ($type_name:ty) => {
  8. paste! {
  9. fn [<read_ $type_name>](&mut self) -> Option<$type_name> {
  10. if self.0.len() == 0 || self.1 + std::mem::size_of::<$type_name>() > self.0.len() {
  11. return None
  12. }
  13. let data = self.read(std::mem::size_of::<$type_name>())?;
  14. unsafe {Some($type_name::from_le(*std::mem::transmute::<_, &$type_name>(data.as_ptr())))}
  15. }
  16. fn [<write_ $type_name>](&mut self, value: $type_name) {
  17. self.write(&value.to_le_bytes());
  18. }
  19. }
  20. };
  21. ($($type_name:ty )+) => {
  22. $(def_rw!($type_name);)+
  23. };
  24. }
  25. macro_rules! def_lmint {
  26. ($name:ident, $type_name:ty) => {
  27. paste! {
  28. fn [<lm_read_ $name>](state: lua_State) -> Result<i32, Box<dyn std::error::Error>> {
  29. unsafe {
  30. let this = &mut *lua::Lcheckudata(state, 1, lua::cstr!("cslb")).cast::<Self>();
  31. match this.[<read_ $type_name>]() {
  32. Some(val) => {
  33. lua::pushinteger(state, val as _);
  34. }
  35. None => {
  36. lua::pushnil(state)
  37. }
  38. }
  39. Ok(1)
  40. }
  41. }
  42. fn [<lm_write_ $name>](state: lua_State) -> Result<i32, Box<dyn std::error::Error>> {
  43. unsafe {
  44. let this = &mut *lua::Lcheckudata(state, 1, lua::cstr!("cslb")).cast::<Self>();
  45. this.[<write_ $type_name>](lua::Lcheckinteger(state, 2) as _);
  46. Ok(0)
  47. }
  48. }
  49. }
  50. };
  51. }
  52. macro_rules! def_lmfloat {
  53. ($name:ident, $type_name:ty) => {
  54. paste! {
  55. fn [<lm_read_ $name>](state: lua_State) -> Result<i32, Box<dyn std::error::Error>> {
  56. unsafe {
  57. let this = &mut *lua::Lcheckudata(state, 1, lua::cstr!("cslb")).cast::<Self>();
  58. match this.[<read_ $type_name>]() {
  59. Some(val) => {
  60. lua::pushnumber(state, val as _);
  61. }
  62. None => {
  63. lua::pushnil(state)
  64. }
  65. }
  66. Ok(1)
  67. }
  68. }
  69. fn [<lm_write_ $name>](state: lua_State) -> Result<i32, Box<dyn std::error::Error>> {
  70. unsafe {
  71. let this = &mut *lua::Lcheckudata(state, 1, lua::cstr!("cslb")).cast::<Self>();
  72. this.[<write_ $type_name>](lua::Lchecknumber(state, 2) as _);
  73. Ok(0)
  74. }
  75. }
  76. }
  77. };
  78. }
  79. pub struct Buffer(Vec<u8>, usize);
  80. impl Buffer {
  81. fn new(size: usize) -> Self {
  82. Self(Vec::with_capacity(size), 0)
  83. }
  84. fn read(&mut self, bytes: usize) -> Option<&[u8]> {
  85. if bytes == 0 || self.0.len() == 0 {
  86. return None;
  87. }
  88. let pos = self.1.min(self.0.len());
  89. let bytes_to_read = if bytes + pos > self.0.len() {
  90. bytes - (bytes + pos - self.0.len())
  91. } else {
  92. bytes
  93. };
  94. if bytes_to_read == 0 {
  95. return None;
  96. }
  97. self.1 += bytes_to_read;
  98. Some(&self.0[pos..pos + bytes_to_read])
  99. }
  100. // Straightforward ripoff from Cursor.
  101. fn write(&mut self, data: &[u8]) {
  102. if self.1 > self.0.len() {
  103. return;
  104. }
  105. let len = self.0.len();
  106. if len < self.1 {
  107. self.0.resize(self.1, 0);
  108. }
  109. {
  110. let space = self.0.len() - self.1;
  111. let (left, right) = data.split_at(std::cmp::min(space, data.len()));
  112. self.0[self.1..self.1 + left.len()].copy_from_slice(left);
  113. self.0.extend_from_slice(right);
  114. }
  115. self.1 += data.len();
  116. }
  117. fn read_f32(&mut self) -> Option<f32> {
  118. if self.0.len() == 0 || self.1 + 4 > self.0.len() {
  119. return None;
  120. }
  121. let data = self.read(4)?;
  122. unsafe {
  123. Some(f32::from_bits(u32::from_le(
  124. *std::mem::transmute::<_, &u32>(data.as_ptr()),
  125. )))
  126. }
  127. }
  128. fn write_f32(&mut self, value: f32) {
  129. self.write(&value.to_le_bytes());
  130. }
  131. fn read_f64(&mut self) -> Option<f64> {
  132. if self.0.len() == 0 || self.1 + 8 > self.0.len() {
  133. return None;
  134. }
  135. let data = self.read(8)?;
  136. unsafe {
  137. Some(f64::from_bits(u64::from_le(
  138. *std::mem::transmute::<_, &u64>(data.as_ptr()),
  139. )))
  140. }
  141. }
  142. fn write_f64(&mut self, value: f64) {
  143. self.write(&value.to_le_bytes());
  144. }
  145. def_rw!(u8 u16 i16 u32 i32);
  146. fn __gc(state: lua_State) -> Result<i32, Box<dyn std::error::Error>> {
  147. unsafe {
  148. let _ = lua::Lcheckudata(state, 1, lua::cstr!("cslb"))
  149. .cast::<Self>()
  150. .read();
  151. Ok(0)
  152. }
  153. }
  154. fn lm_read(state: lua_State) -> Result<i32, Box<dyn std::error::Error>> {
  155. unsafe {
  156. let this = &mut *lua::Lcheckudata(state, 1, lua::cstr!("cslb")).cast::<Self>();
  157. match this.read(lua::Lcheckinteger(state, 2) as _) {
  158. Some(val) => {
  159. lua::pushlstring(state, val.as_ptr(), val.len());
  160. }
  161. None => lua::pushnil(state),
  162. }
  163. Ok(1)
  164. }
  165. }
  166. fn lm_write(state: lua_State) -> Result<i32, Box<dyn std::error::Error>> {
  167. unsafe {
  168. let this = &mut *lua::Lcheckudata(state, 1, lua::cstr!("cslb")).cast::<Self>();
  169. this.write(check_slice!(state, 2));
  170. Ok(0)
  171. }
  172. }
  173. fn lm_read_bool(state: lua_State) -> Result<i32, Box<dyn std::error::Error>> {
  174. unsafe {
  175. let this = &mut *lua::Lcheckudata(state, 1, lua::cstr!("cslb")).cast::<Self>();
  176. match this.read_u8() {
  177. Some(val) => {
  178. lua::pushboolean(state, val as _);
  179. }
  180. None => lua::pushnil(state),
  181. }
  182. Ok(1)
  183. }
  184. }
  185. fn lm_write_bool(state: lua_State) -> Result<i32, Box<dyn std::error::Error>> {
  186. unsafe {
  187. let this = &mut *lua::Lcheckudata(state, 1, lua::cstr!("cslb")).cast::<Self>();
  188. this.write_u8(lua::toboolean(state, 2) as _);
  189. Ok(0)
  190. }
  191. }
  192. def_lmint!(byte, u8);
  193. def_lmint!(ushort, u16);
  194. def_lmint!(short, i16);
  195. def_lmint!(ulong, u32);
  196. def_lmint!(long, i32);
  197. def_lmfloat!(float, f32);
  198. def_lmfloat!(double, f64);
  199. fn lm_tell(state: lua_State) -> Result<i32, Box<dyn std::error::Error>> {
  200. unsafe {
  201. let this = &mut *lua::Lcheckudata(state, 1, lua::cstr!("cslb")).cast::<Self>();
  202. lua::pushinteger(state, this.1 as _);
  203. Ok(1)
  204. }
  205. }
  206. fn lm_seek(state: lua_State) -> Result<i32, Box<dyn std::error::Error>> {
  207. unsafe {
  208. let this = &mut *lua::Lcheckudata(state, 1, lua::cstr!("cslb")).cast::<Self>();
  209. this.1 = (lua::Lcheckinteger(state, 2).max(0) as usize).min(this.0.len() - 1);
  210. Ok(0)
  211. }
  212. }
  213. fn lm_clear(state: lua_State) -> Result<i32, Box<dyn std::error::Error>> {
  214. unsafe {
  215. let this = &mut *lua::Lcheckudata(state, 1, lua::cstr!("cslb")).cast::<Self>();
  216. this.0.clear();
  217. this.1 = 0;
  218. Ok(0)
  219. }
  220. }
  221. fn lm_resize(state: lua_State) -> Result<i32, Box<dyn std::error::Error>> {
  222. unsafe {
  223. let this = &mut *lua::Lcheckudata(state, 1, lua::cstr!("cslb")).cast::<Self>();
  224. let size = lua::Lcheckinteger(state, 2).max(0) as usize;
  225. this.0.resize(size, 0);
  226. this.1 = this.1.min(size);
  227. Ok(0)
  228. }
  229. }
  230. fn lm_shrink(state: lua_State) -> Result<i32, Box<dyn std::error::Error>> {
  231. unsafe {
  232. let this = &mut *lua::Lcheckudata(state, 1, lua::cstr!("cslb")).cast::<Self>();
  233. this.0.shrink_to_fit();
  234. Ok(1)
  235. }
  236. }
  237. fn lm_get_value(state: lua_State) -> Result<i32, Box<dyn std::error::Error>> {
  238. unsafe {
  239. let this = &mut *lua::Lcheckudata(state, 1, lua::cstr!("cslb")).cast::<Self>();
  240. lua::pushlstring(state, this.0.as_ptr(), this.0.len());
  241. Ok(1)
  242. }
  243. }
  244. fn lm_set_value(state: lua_State) -> Result<i32, Box<dyn std::error::Error>> {
  245. unsafe {
  246. let this = &mut *lua::Lcheckudata(state, 1, lua::cstr!("cslb")).cast::<Self>();
  247. let data = check_slice!(state, 2);
  248. this.0.clear();
  249. this.1 = 0;
  250. this.0.extend_from_slice(data);
  251. Ok(0)
  252. }
  253. }
  254. fn metatable(state: lua_State) {
  255. unsafe {
  256. if lua::Lnewmetatable(state, lua::cstr!("cslb")) {
  257. lua::pushvalue(state, -1);
  258. lua::setfield(state, -2, lua::cstr!("__index"));
  259. insert_function!(state, "__gc", Self::__gc);
  260. insert_function!(state, "Read", Self::lm_read);
  261. insert_function!(state, "ReadBool", Self::lm_read_bool);
  262. insert_function!(state, "ReadByte", Self::lm_read_byte);
  263. insert_function!(state, "ReadUShort", Self::lm_read_ushort);
  264. insert_function!(state, "ReadShort", Self::lm_read_short);
  265. insert_function!(state, "ReadULong", Self::lm_read_ulong);
  266. insert_function!(state, "ReadLong", Self::lm_read_long);
  267. insert_function!(state, "ReadFloat", Self::lm_read_float);
  268. insert_function!(state, "ReadDouble", Self::lm_read_double);
  269. insert_function!(state, "Write", Self::lm_write);
  270. insert_function!(state, "WriteBool", Self::lm_write_bool);
  271. insert_function!(state, "WriteByte", Self::lm_write_byte);
  272. insert_function!(state, "WriteUShort", Self::lm_write_ushort);
  273. insert_function!(state, "WriteShort", Self::lm_write_short);
  274. insert_function!(state, "WriteULong", Self::lm_write_ulong);
  275. insert_function!(state, "WriteLong", Self::lm_write_long);
  276. insert_function!(state, "WriteFloat", Self::lm_write_float);
  277. insert_function!(state, "WriteDouble", Self::lm_write_double);
  278. insert_function!(state, "Tell", Self::lm_tell);
  279. insert_function!(state, "Seek", Self::lm_seek);
  280. insert_function!(state, "Clear", Self::lm_clear);
  281. insert_function!(state, "Resize", Self::lm_resize);
  282. insert_function!(state, "Shrink", Self::lm_shrink);
  283. insert_function!(state, "GetValue", Self::lm_get_value);
  284. insert_function!(state, "SetValue", Self::lm_set_value);
  285. }
  286. }
  287. }
  288. pub fn l_new(state: lua_State) -> Result<i32, Box<dyn std::error::Error>> {
  289. unsafe {
  290. let buffer = match lua::get_type(state, 1) {
  291. 3 => Self::new(lua::tonumber(state, 1) as _),
  292. 4 => {
  293. let data = check_slice!(state, 1);
  294. let mut buffer = Self::new(data.len());
  295. buffer.0.extend_from_slice(data);
  296. buffer
  297. },
  298. _ => {Self::new(0)}
  299. };
  300. let udata = lua::newuserdata(state, std::mem::size_of::<Self>()).cast::<Self>();
  301. udata.write(buffer);
  302. Self::metatable(state);
  303. lua::setmetatable(state, -2);
  304. }
  305. Ok(1)
  306. }
  307. }