lua-range/_range.lua at master · benpop/lua-range · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
local _M = {
  _AUTHORS = "benpop",
  _VERSION = "1.2",
  _DESCRIPTION = "range generator (cf. Python (x)range, Ruby \"..\" operator)"
}


local _mt = {}  -- metatable


-- see Python-2.7.3/Objects/rangeobject.c:19
local function _len (lo, hi, step)
  local n
  if step > 0 and lo <= hi then
    n = 1 + (hi - lo) / step
  elseif step < 0 and lo >= hi then
    n = 1 + (lo - hi) / -step
  else
    n = 0
  end
  return math.floor(n)
end


local function toint (x)
  local n = tonumber(x)
  if n == nil then
    return nil
  end
  -- truncate toward zero
  if n > 0 then
    return math.floor(n)
  elseif n < 0 then
    return math.ceil(n)
  else
    return 0
  end
end


local function checkint (name, argn, ...)
  local arg = select(argn, ...)
  return toint(arg) or
      error(string.format(
          "bad argument #%d: '%s' (expected integer, got %s)",
          argn, name, type(arg)), 2)
end


function _M.new (...)
  local n = select("#", ...)
  local lo, hi, step
  if n == 1 then
    hi = checkint("stop", 1, ...)
    assert(hi >= 1, "single argument 'stop' must be positive")
    lo, step = 1, 1
  elseif n == 2 then
    lo = checkint("start", 1, ...)
    hi = checkint("stop", 2, ...)
    step = lo <= hi and 1 or -1
  elseif n == 3 then
    lo = checkint("start", 1, ...)
    hi = checkint("stop", 2, ...)
    step = checkint("step", 3, ...)
  else
    error("needs 1-3 arguments")
  end
  return setmetatable({
        start = lo,
        stop = hi,
        step = step,
        len = _len(lo, hi, step)
      }, _mt)
end


function _M.bless (self)
  if type(self) ~= "table" then
    return nil, "not a table"
  elseif not self.stop then
    return nil, "no valid fields"
  end
  if getmetatable(self) == _mt and self.start and self.stop and
      self.step and self.len then
    return self
  end
  if not self.start and self.stop < 0 then
    error("single field 'stop' must be positive")
  end
  if not self.start then self.start = 1 end
  if not self.step then self.step = self.start <= self.stop and 1 or -1 end
  self.len = _len(self.start, self.stop, self.step)
  return setmetatable(self, _mt)
end


function _M:__call (...)
  return self.new(...)
end


function _mt:get (i)
  if 1 <= i and i <= self.len then
    return (i - 1) * self.step + self.start
  end
end


function _mt:__tostring ()
  if self.start == 1 and self.step == 1 then
    return string.format("range(%d)", self.stop)
  elseif (self.step == 1 and self.start <= self.stop) or
      (self.step == -1 and self.start >= self.stop) then
    return string.format("range(%d, %d)", self.start, self.stop)
  else
    return string.format("range(%d, %d, %d)",
        self.start, self.stop, self.step)
  end
end


function _mt:__len ()
  return self.len
end


function _mt:__index (k)
  local i = tonumber(k)
  if i then
    return self:get(i)
  else
    return _mt[k]
  end
end


local function _next (self, i)
  i = i + 1  -- next index
  local v = self:get(i)
  if v then return i, v end
end


function _mt:__ipairs ()
  return _next, self, 0
end


function _mt:totable ()
  local t = {}
  for i,v in ipairs(self) do
    t[i] = v
  end
  return t
end


return setmetatable(_M, _M)