trickle.lua 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222
  1. local trickle = {}
  2. local function byteExtract(x, a, b)
  3. b = b or a
  4. x = x % (2 ^ (b + 1))
  5. for i = 1, a do
  6. x = math.floor(x / 2)
  7. end
  8. return x
  9. end
  10. local function byteInsert(x, y, a, b)
  11. local res = x
  12. for i = a, b do
  13. local e = byteExtract(y, i - a)
  14. if e ~= byteExtract(x, i) then
  15. res = (e == 1) and res + (2 ^ i) or res - (2 ^ i)
  16. end
  17. end
  18. return res
  19. end
  20. function trickle.create(str)
  21. local stream = {
  22. str = str or '',
  23. byte = nil,
  24. byteLen = nil
  25. }
  26. return setmetatable(stream, trickle)
  27. end
  28. function trickle:truncate()
  29. if self.byte then
  30. self.str = self.str .. string.char(self.byte)
  31. self.byte = nil
  32. self.byteLen = nil
  33. end
  34. return self.str
  35. end
  36. function trickle:clear()
  37. self.str = ''
  38. self.byte = nil
  39. self.byteLen = nil
  40. return self
  41. end
  42. function trickle:write(x, sig)
  43. if sig == 'string' then self:writeString(x)
  44. elseif sig == 'bool' then self:writeBool(x)
  45. elseif sig == 'float' then self:writeFloat(x)
  46. else
  47. local n = sig:match('(%d+)bit')
  48. self:writeBits(x, n)
  49. end
  50. return self
  51. end
  52. function trickle:writeString(string)
  53. self:truncate()
  54. string = tostring(string)
  55. self.str = self.str .. string.char(#string) .. string
  56. end
  57. function trickle:writeBool(bool)
  58. local x = bool and 1 or 0
  59. self:writeBits(x, 1)
  60. end
  61. function trickle:writeFloat(float)
  62. self:writeString(float)
  63. end
  64. function trickle:writeBits(x, n)
  65. local idx = 0
  66. repeat
  67. if not self.byte then self.byte = 0 self.byteLen = 0 end
  68. local numWrite = math.min(n, (7 - self.byteLen) + 1)
  69. local toWrite = byteExtract(x, idx, idx + (numWrite - 1))
  70. self.byte = byteInsert(self.byte, toWrite, self.byteLen, self.byteLen + (numWrite - 1))
  71. self.byteLen = self.byteLen + numWrite
  72. if self.byteLen == 8 then
  73. self.str = self.str .. string.char(self.byte)
  74. self.byte = nil
  75. self.byteLen = nil
  76. end
  77. n = n - numWrite
  78. idx = idx + numWrite
  79. until n == 0
  80. end
  81. function trickle:read(kind)
  82. if kind == 'string' then return self:readString()
  83. elseif kind == 'bool' then return self:readBool()
  84. elseif kind == 'float' then return self:readFloat()
  85. else
  86. local n = tonumber(kind:match('(%d+)bit'))
  87. return self:readBits(n)
  88. end
  89. end
  90. function trickle:readString()
  91. if self.byte then
  92. self.str = self.str:sub(2)
  93. self.byte = nil
  94. self.byteLen = nil
  95. end
  96. local len = self.str:byte(1)
  97. local res = ''
  98. if len then
  99. self.str = self.str:sub(2)
  100. res = self.str:sub(1, len)
  101. self.str = self.str:sub(len + 1)
  102. end
  103. return res
  104. end
  105. function trickle:readBool()
  106. return self:readBits(1) > 0
  107. end
  108. function trickle:readFloat()
  109. return tonumber(self:readString())
  110. end
  111. function trickle:readBits(n)
  112. local x = 0
  113. local idx = 0
  114. while n > 0 do
  115. if not self.byte then self.byte = self.str:byte(1) or 0 self.byteLen = 0 end
  116. local numRead = math.min(n, (7 - self.byteLen) + 1)
  117. x = x + (byteExtract(self.byte, self.byteLen, self.byteLen + (numRead - 1)) * (2 ^ idx))
  118. self.byteLen = self.byteLen + numRead
  119. if self.byteLen == 8 then
  120. self.str = self.str:sub(2)
  121. self.byte = nil
  122. self.byteLen = nil
  123. end
  124. n = n - numRead
  125. idx = idx + numRead
  126. end
  127. return x
  128. end
  129. function trickle:pack(data, signature)
  130. local keys
  131. if signature.delta then
  132. keys = {}
  133. for _, key in ipairs(signature.delta) do
  134. if type(key) == 'table' then
  135. local has = 0
  136. for i = 1, #key do
  137. if data[key[i]] ~= nil then
  138. keys[key[i]] = true
  139. has = has + 1
  140. else
  141. keys[key[i]] = false
  142. end
  143. end
  144. if has == 0 then self:write(0, '1bit')
  145. elseif has == #key then self:write(1, '1bit')
  146. else error('Only part of message delta group "' .. table.concat(key, ', ') .. '" was provided.') end
  147. else
  148. self:write(data[key] ~= nil and 1 or 0, '1bit')
  149. keys[key] = data[key] ~= nil and true or false
  150. end
  151. end
  152. end
  153. for _, sig in ipairs(signature) do
  154. if not keys or keys[sig[1]] ~= false then
  155. if type(sig[2]) == 'table' then
  156. self:write(#data[sig[1]], '4bits')
  157. for i = 1, #data[sig[1]] do self:pack(data[sig[1]][i], sig[2]) end
  158. else
  159. self:write(data[sig[1]], sig[2])
  160. end
  161. end
  162. end
  163. end
  164. function trickle:unpack(signature)
  165. local keys
  166. if signature.delta then
  167. keys = {}
  168. for i = 1, #signature.delta do
  169. local val = self:read('1bit') > 0
  170. if type(signature.delta[i]) == 'table' then
  171. for j = 1, #signature.delta[i] do keys[signature.delta[i][j]] = val end
  172. else
  173. keys[signature.delta[i]] = val
  174. end
  175. end
  176. end
  177. local data = {}
  178. for _, sig in ipairs(signature) do
  179. if not keys or keys[sig[1]] ~= false then
  180. if type(sig[2]) == 'table' then
  181. local ct = self:read('4bits')
  182. data[sig[1]] = {}
  183. for i = 1, ct do table.insert(data[sig[1]], self:unpack(sig[2])) end
  184. else
  185. data[sig[1]] = self:read(sig[2])
  186. end
  187. end
  188. end
  189. return data
  190. end
  191. trickle.__tostring = trickle.truncate
  192. trickle.__index = trickle
  193. return { create = trickle.create }