github.com/xrash/gopher-lua@v0.0.0-20160304065408-e5faab4db06a/stringlib.go (about)

     1  package lua
     2  
     3  import (
     4  	"fmt"
     5  	"strings"
     6  	"unsafe"
     7  
     8  	"github.com/yuin/gopher-lua/pm"
     9  )
    10  
    11  func OpenString(L *LState) int {
    12  	var mod *LTable
    13  	//_, ok := L.G.builtinMts[int(LTString)]
    14  	//if !ok {
    15  	mod = L.RegisterModule(StringLibName, strFuncs).(*LTable)
    16  	gmatch := L.NewClosure(strGmatch, L.NewFunction(strGmatchIter))
    17  	mod.RawSetString("gmatch", gmatch)
    18  	mod.RawSetString("gfind", gmatch)
    19  	mod.RawSetString("__index", mod)
    20  	L.G.builtinMts[int(LTString)] = mod
    21  	//}
    22  	L.Push(mod)
    23  	return 1
    24  }
    25  
    26  var strFuncs = map[string]LGFunction{
    27  	"byte":    strByte,
    28  	"char":    strChar,
    29  	"dump":    strDump,
    30  	"find":    strFind,
    31  	"format":  strFormat,
    32  	"gsub":    strGsub,
    33  	"len":     strLen,
    34  	"lower":   strLower,
    35  	"match":   strMatch,
    36  	"rep":     strRep,
    37  	"reverse": strReverse,
    38  	"sub":     strSub,
    39  	"upper":   strUpper,
    40  }
    41  
    42  func strByte(L *LState) int {
    43  	str := L.CheckString(1)
    44  	start := L.OptInt(2, 1) - 1
    45  	end := L.OptInt(3, -1)
    46  	l := len(str)
    47  	if start < 0 {
    48  		start = l + start + 1
    49  	}
    50  	if end < 0 {
    51  		end = l + end + 1
    52  	}
    53  
    54  	if L.GetTop() == 2 {
    55  		if start < 0 || start >= l {
    56  			return 0
    57  		}
    58  		L.Push(LNumber(str[start]))
    59  		return 1
    60  	}
    61  
    62  	start = intMax(start, 0)
    63  	end = intMin(end, l)
    64  	if end < 0 || end <= start || start >= l {
    65  		return 0
    66  	}
    67  
    68  	for i := start; i < end; i++ {
    69  		L.Push(LNumber(str[i]))
    70  	}
    71  	return end - start
    72  }
    73  
    74  func strChar(L *LState) int {
    75  	top := L.GetTop()
    76  	bytes := make([]byte, L.GetTop())
    77  	for i := 1; i <= top; i++ {
    78  		bytes[i-1] = uint8(L.CheckInt(i))
    79  	}
    80  	L.Push(LString(string(bytes)))
    81  	return 1
    82  }
    83  
    84  func strDump(L *LState) int {
    85  	L.RaiseError("GopherLua does not support the string.dump")
    86  	return 0
    87  }
    88  
    89  func strFind(L *LState) int {
    90  	str := L.CheckString(1)
    91  	pattern := L.CheckString(2)
    92  	if len(pattern) == 0 {
    93  		L.Push(LNumber(1))
    94  		L.Push(LNumber(0))
    95  		return 2
    96  	}
    97  	init := luaIndex2StringIndex(str, L.OptInt(3, 1), true)
    98  	plain := false
    99  	if L.GetTop() == 4 {
   100  		plain = LVAsBool(L.Get(4))
   101  	}
   102  	if len(str) == 0 && len(pattern) == 0 {
   103  		L.Push(LNumber(1))
   104  		return 1
   105  	}
   106  
   107  	if plain {
   108  		pos := strings.Index(str[init:], pattern)
   109  		if pos < 0 {
   110  			L.Push(LNil)
   111  			return 1
   112  		}
   113  		L.Push(LNumber(init+pos) + 1)
   114  		L.Push(LNumber(init + pos + len(pattern)))
   115  		return 2
   116  	}
   117  
   118  	mds, err := pm.Find(pattern, *(*[]byte)(unsafe.Pointer(&str)), init, 1)
   119  	if err != nil {
   120  		L.RaiseError(err.Error())
   121  	}
   122  	if len(mds) == 0 {
   123  		L.Push(LNil)
   124  		return 1
   125  	}
   126  	md := mds[0]
   127  	L.Push(LNumber(md.Capture(0) + 1))
   128  	L.Push(LNumber(md.Capture(1)))
   129  	for i := 2; i < md.CaptureLength(); i += 2 {
   130  		if md.IsPosCapture(i) {
   131  			L.Push(LNumber(md.Capture(i)))
   132  		} else {
   133  			L.Push(LString(str[md.Capture(i):md.Capture(i+1)]))
   134  		}
   135  	}
   136  	return md.CaptureLength()/2 + 1
   137  }
   138  
   139  func strFormat(L *LState) int {
   140  	str := L.CheckString(1)
   141  	args := make([]interface{}, L.GetTop()-1)
   142  	top := L.GetTop()
   143  	for i := 2; i <= top; i++ {
   144  		args[i-2] = L.Get(i)
   145  	}
   146  	npat := strings.Count(str, "%") - strings.Count(str, "%%")
   147  	L.Push(LString(fmt.Sprintf(str, args[:intMin(npat, len(args))]...)))
   148  	return 1
   149  }
   150  
   151  func strGsub(L *LState) int {
   152  	str := L.CheckString(1)
   153  	pat := L.CheckString(2)
   154  	L.CheckTypes(3, LTString, LTTable, LTFunction)
   155  	repl := L.CheckAny(3)
   156  	limit := L.OptInt(4, -1)
   157  
   158  	mds, err := pm.Find(pat, *(*[]byte)(unsafe.Pointer(&str)), 0, limit)
   159  	if err != nil {
   160  		L.RaiseError(err.Error())
   161  	}
   162  	if len(mds) == 0 {
   163  		L.SetTop(1)
   164  		L.Push(LNumber(0))
   165  		return 2
   166  	}
   167  	switch lv := repl.(type) {
   168  	case LString:
   169  		L.Push(LString(strGsubStr(L, str, string(lv), mds)))
   170  	case *LTable:
   171  		L.Push(LString(strGsubTable(L, str, lv, mds)))
   172  	case *LFunction:
   173  		L.Push(LString(strGsubFunc(L, str, lv, mds)))
   174  	}
   175  	L.Push(LNumber(len(mds)))
   176  	return 2
   177  }
   178  
   179  type replaceInfo struct {
   180  	Indicies []int
   181  	String   string
   182  }
   183  
   184  func checkCaptureIndex(L *LState, m *pm.MatchData, idx int) {
   185  	if idx <= 2 {
   186  		return
   187  	}
   188  	if idx >= m.CaptureLength() {
   189  		L.RaiseError("invalid capture index")
   190  	}
   191  }
   192  
   193  func capturedString(L *LState, m *pm.MatchData, str string, idx int) string {
   194  	checkCaptureIndex(L, m, idx)
   195  	if idx >= m.CaptureLength() && idx == 2 {
   196  		idx = 0
   197  	}
   198  	if m.IsPosCapture(idx) {
   199  		return fmt.Sprint(m.Capture(idx))
   200  	} else {
   201  		return str[m.Capture(idx):m.Capture(idx+1)]
   202  	}
   203  
   204  }
   205  
   206  func strGsubDoReplace(str string, info []replaceInfo) string {
   207  	offset := 0
   208  	buf := []byte(str)
   209  	for _, replace := range info {
   210  		oldlen := len(buf)
   211  		b1 := append([]byte(""), buf[0:offset+replace.Indicies[0]]...)
   212  		b2 := []byte("")
   213  		index2 := offset + replace.Indicies[1]
   214  		if index2 <= len(buf) {
   215  			b2 = append(b2, buf[index2:len(buf)]...)
   216  		}
   217  		buf = append(b1, replace.String...)
   218  		buf = append(buf, b2...)
   219  		offset += len(buf) - oldlen
   220  	}
   221  	return string(buf)
   222  }
   223  
   224  func strGsubStr(L *LState, str string, repl string, matches []*pm.MatchData) string {
   225  	infoList := make([]replaceInfo, 0, len(matches))
   226  	for _, match := range matches {
   227  		start, end := match.Capture(0), match.Capture(1)
   228  		sc := newFlagScanner('%', "", "", repl)
   229  		for c, eos := sc.Next(); !eos; c, eos = sc.Next() {
   230  			if !sc.ChangeFlag {
   231  				if sc.HasFlag {
   232  					if c >= '0' && c <= '9' {
   233  						sc.AppendString(capturedString(L, match, str, 2*(int(c)-48)))
   234  					} else {
   235  						sc.AppendChar('%')
   236  						sc.AppendChar(c)
   237  					}
   238  					sc.HasFlag = false
   239  				} else {
   240  					sc.AppendChar(c)
   241  				}
   242  			}
   243  		}
   244  		infoList = append(infoList, replaceInfo{[]int{start, end}, sc.String()})
   245  	}
   246  
   247  	return strGsubDoReplace(str, infoList)
   248  }
   249  
   250  func strGsubTable(L *LState, str string, repl *LTable, matches []*pm.MatchData) string {
   251  	infoList := make([]replaceInfo, 0, len(matches))
   252  	for _, match := range matches {
   253  		idx := 0
   254  		if match.CaptureLength() > 2 { // has captures
   255  			idx = 2
   256  		}
   257  		var value LValue
   258  		if match.IsPosCapture(idx) {
   259  			value = L.GetTable(repl, LNumber(match.Capture(idx)))
   260  		} else {
   261  			value = L.GetField(repl, str[match.Capture(idx):match.Capture(idx+1)])
   262  		}
   263  		if !LVIsFalse(value) {
   264  			infoList = append(infoList, replaceInfo{[]int{match.Capture(0), match.Capture(1)}, LVAsString(value)})
   265  		}
   266  	}
   267  	return strGsubDoReplace(str, infoList)
   268  }
   269  
   270  func strGsubFunc(L *LState, str string, repl *LFunction, matches []*pm.MatchData) string {
   271  	infoList := make([]replaceInfo, 0, len(matches))
   272  	for _, match := range matches {
   273  		start, end := match.Capture(0), match.Capture(1)
   274  		L.Push(repl)
   275  		nargs := 0
   276  		if match.CaptureLength() > 2 { // has captures
   277  			for i := 2; i < match.CaptureLength(); i += 2 {
   278  				if match.IsPosCapture(i) {
   279  					L.Push(LNumber(match.Capture(i)))
   280  				} else {
   281  					L.Push(LString(capturedString(L, match, str, i)))
   282  				}
   283  				nargs++
   284  			}
   285  		} else {
   286  			L.Push(LString(capturedString(L, match, str, 0)))
   287  			nargs++
   288  		}
   289  		L.Call(nargs, 1)
   290  		ret := L.reg.Pop()
   291  		if !LVIsFalse(ret) {
   292  			infoList = append(infoList, replaceInfo{[]int{start, end}, LVAsString(ret)})
   293  		}
   294  	}
   295  	return strGsubDoReplace(str, infoList)
   296  }
   297  
   298  type strMatchData struct {
   299  	str     string
   300  	pos     int
   301  	matches []*pm.MatchData
   302  }
   303  
   304  func strGmatchIter(L *LState) int {
   305  	md := L.CheckUserData(1).Value.(*strMatchData)
   306  	str := md.str
   307  	matches := md.matches
   308  	idx := md.pos
   309  	md.pos += 1
   310  	if idx == len(matches) {
   311  		return 0
   312  	}
   313  	L.Push(L.Get(1))
   314  	match := matches[idx]
   315  	if match.CaptureLength() == 2 {
   316  		L.Push(LString(str[match.Capture(0):match.Capture(1)]))
   317  		return 1
   318  	}
   319  
   320  	for i := 2; i < match.CaptureLength(); i += 2 {
   321  		if match.IsPosCapture(i) {
   322  			L.Push(LNumber(match.Capture(i)))
   323  		} else {
   324  			L.Push(LString(str[match.Capture(i):match.Capture(i+1)]))
   325  		}
   326  	}
   327  	return match.CaptureLength()/2 - 1
   328  }
   329  
   330  func strGmatch(L *LState) int {
   331  	str := L.CheckString(1)
   332  	pattern := L.CheckString(2)
   333  	mds, err := pm.Find(pattern, []byte(str), 0, -1)
   334  	if err != nil {
   335  		L.RaiseError(err.Error())
   336  	}
   337  	L.Push(L.Get(UpvalueIndex(1)))
   338  	ud := L.NewUserData()
   339  	ud.Value = &strMatchData{str, 0, mds}
   340  	L.Push(ud)
   341  	return 2
   342  }
   343  
   344  func strLen(L *LState) int {
   345  	str := L.CheckString(1)
   346  	L.Push(LNumber(len(str)))
   347  	return 1
   348  }
   349  
   350  func strLower(L *LState) int {
   351  	str := L.CheckString(1)
   352  	L.Push(LString(strings.ToLower(str)))
   353  	return 1
   354  }
   355  
   356  func strMatch(L *LState) int {
   357  	str := L.CheckString(1)
   358  	pattern := L.CheckString(2)
   359  	offset := L.OptInt(3, 1)
   360  	l := len(str)
   361  	if offset < 0 {
   362  		offset = l + offset + 1
   363  	}
   364  	offset--
   365  	if offset < 0 {
   366  		offset = 0
   367  	}
   368  
   369  	mds, err := pm.Find(pattern, *(*[]byte)(unsafe.Pointer(&str)), offset, 1)
   370  	if err != nil {
   371  		L.RaiseError(err.Error())
   372  	}
   373  	if len(mds) == 0 {
   374  		L.Push(LNil)
   375  		return 0
   376  	}
   377  	md := mds[0]
   378  	nsubs := md.CaptureLength() / 2
   379  	switch nsubs {
   380  	case 1:
   381  		L.Push(LString(str[md.Capture(0):md.Capture(1)]))
   382  		return 1
   383  	default:
   384  		for i := 2; i < md.CaptureLength(); i += 2 {
   385  			if md.IsPosCapture(i) {
   386  				L.Push(LNumber(md.Capture(i)))
   387  			} else {
   388  				L.Push(LString(str[md.Capture(i):md.Capture(i+1)]))
   389  			}
   390  		}
   391  		return nsubs - 1
   392  	}
   393  }
   394  
   395  func strRep(L *LState) int {
   396  	str := L.CheckString(1)
   397  	n := L.CheckInt(2)
   398  	L.Push(LString(strings.Repeat(str, n)))
   399  	return 1
   400  }
   401  
   402  func strReverse(L *LState) int {
   403  	str := L.CheckString(1)
   404  	bts := []byte(str)
   405  	out := make([]byte, len(bts))
   406  	for i, j := 0, len(bts)-1; j >= 0; i, j = i+1, j-1 {
   407  		out[i] = bts[j]
   408  	}
   409  	L.Push(LString(string(out)))
   410  	return 1
   411  }
   412  
   413  func strSub(L *LState) int {
   414  	str := L.CheckString(1)
   415  	start := luaIndex2StringIndex(str, L.CheckInt(2), true)
   416  	end := luaIndex2StringIndex(str, L.OptInt(3, -1), false)
   417  	l := len(str)
   418  	if start >= l || end < start {
   419  		L.Push(LString(""))
   420  	} else {
   421  		L.Push(LString(str[start:end]))
   422  	}
   423  	return 1
   424  }
   425  
   426  func strUpper(L *LState) int {
   427  	str := L.CheckString(1)
   428  	L.Push(LString(strings.ToUpper(str)))
   429  	return 1
   430  }
   431  
   432  func luaIndex2StringIndex(str string, i int, start bool) int {
   433  	if start && i != 0 {
   434  		i -= 1
   435  	}
   436  	l := len(str)
   437  	if i < 0 {
   438  		i = l + i + 1
   439  	}
   440  	i = intMax(0, i)
   441  	if !start && i > l {
   442  		i = l
   443  	}
   444  	return i
   445  }
   446  
   447  //