def _EncodeAny(t, f, dl, s, issueError):
s = ut.unProxy(s)
if isinstance(s, str) or isinstance(s, bytes):
try:
return f(s)
except KeyError:
if isinstance(s, str):
try:
return dl[s]
except KeyError:
sl = ["Cannot encode string %s as %s." % (s, t)]
for match in difflib.get_close_matches(s, dl.keys()):
sl.append(' - Suggestion: %s' % match)
raise ut.EPError('\n'.join(sl))
if issueError:
raise ut.EPError('[Warning] "%s" is not a %s' % (s, t))
return s
else:
try:
return dl.get(s, s)
except TypeError: # unhashable
return s
def AddString(self, string):
string = ut.u2b(string) # Starcraft uses multibyte encoding.
if not isinstance(string, bytes):
raise ut.EPError('Invalid type for string')
stringindex = len(self._dataindextb)
# If duplicate text exist -> just proxy it
try:
repr_stringid = self._stringmap[string]
dataindex = self._dataindextb[repr_stringid]
self._dataindextb.append(dataindex)
self._capacity += 2 # just string offset
# Else -> Create new entry
except KeyError:
dataindex = len(self._datatb)
self._stringmap[string] = stringindex
self._datatb.append(string)
self._dataindextb.append(dataindex)
# string + b'\0' + string offset
self._capacity += len(string) + 1 + 2
ut.ep_assert(self._capacity < 65536, 'String table overflow')
return stringindex
def CollectObjects(root):
global phase
global _rootobj
global _found_objects
global _found_objects_set
global _dynamic_objects_set
global _untraversed_objects
lprint('[Stage 1/3] CollectObjects', flush=True)
phase = PHASE_COLLECTING
objc = ObjCollector()
_rootobj = None
_found_objects_set = set()
_dynamic_objects_set = set()
_untraversed_objects = []
# Evaluate root to register root object.
# root may not have WritePayload() method e.g: Forward()
constexpr.Evaluate(root)
while _untraversed_objects:
while _untraversed_objects:
lprint(" - Collected %d / %d objects" %
(len(_found_objects_set),
len(_found_objects_set) + len(_untraversed_objects)))
obj = _untraversed_objects.pop()
objc.StartWrite()
obj.CollectDependency(objc)
objc.EndWrite()
# Check for new objects
for obj in list(_dynamic_objects_set):
objc.StartWrite()
obj.CollectDependency(objc)
objc.EndWrite()
if len(_found_objects_set) == 0:
raise ut.EPError('No object collected')
# Shuffle objects -> Randomize(?) addresses
_found_objects_set.remove(_rootobj)
fo2 = list(_found_objects_set)
random.shuffle(fo2)
_found_objects = [_rootobj] + fo2
# cleanup
_found_objects_set = None
phase = None
# Final
lprint(" - Collected %d / %d objects" %
(len(_found_objects), len(_found_objects)),
flush=True)
def GetStringIndex(self, string):
string = ut.u2b(string)
if not isinstance(string, bytes):
raise ut.EPError('Invalid type for string')
try:
return self._stringmap[string] + 1
except KeyError:
return self.AddString(string) + 1
def AllocObjects():
global phase
global _alloctable
global _payload_size
phase = PHASE_ALLOCATING
objn = len(_found_objects)
lprint("[Stage 2/3] AllocObjects", flush=True)
# Quick and less space-efficient approach
if not _payload_compress:
lallocaddr = 0
for i, obj in enumerate(_found_objects):
objsize = obj.GetDataSize()
allocsize = (objsize + 3) & ~3
_alloctable[obj] = lallocaddr
lallocaddr += allocsize
lprint(" - Allocated %d / %d objects" % (i + 1, objn))
_payload_size = lallocaddr
lprint(" - Allocated %d / %d objects" % (objn, objn), flush=True)
phase = None
return
obja = ObjAllocator()
_alloctable = {}
dwoccupmap_dict = {}
# Get occupation map for all objects
for i, obj in enumerate(_found_objects):
obja.StartWrite()
obj.WritePayload(obja)
dwoccupmap = obja.EndWrite()
dwoccupmap_dict[obj] = dwoccupmap
if len(dwoccupmap) != (obj.GetDataSize() + 3) >> 2:
raise ut.EPError(
'Occupation map length (%d) & Object size mismatch for object (%d)'
% (len(dwoccupmap), (obj.GetDataSize() + 3) >> 2))
lprint(" - Preprocessed %d / %d objects" % (i + 1, objn))
lprint(" - Preprocessed %d / %d objects" % (objn, objn), flush=True)
lprint(" - Allocating objects..", flush=True)
stackobjs.StackObjects(_found_objects, dwoccupmap_dict, _alloctable)
# Get payload length
_payload_size = max(
map(lambda obj: _alloctable[obj] + obj.GetDataSize(), _found_objects))
phase = None
def CompressPayload(mode):
''' Set payload compression mode.
:param mode: If true, enable object stacking (compression). If false,
disable it.
'''
global _payload_compress
if mode not in [True, False]:
raise ut.EPError('Invalid type')
if mode:
_payload_compress = True
else:
_payload_compress = False
def WriteDword(self, number):
number = constexpr.Evaluate(number)
if number.rlocmode:
ut.ep_assert(
self._datacur % 4 == 0,
'Non-const dwords must be aligned to 4byte'
)
if number.rlocmode == 1:
self._prttable.append(self._datacur)
elif number.rlocmode == 4:
self._orttable.append(self._datacur)
else:
raise ut.EPError('rlocmode should be 1 or 4')
offset = number.offset
self._data[self._datacur + 0] = offset & 0xFF
self._data[self._datacur + 1] = (offset >> 8) & 0xFF
self._data[self._datacur + 2] = (offset >> 16) & 0xFF
self._data[self._datacur + 3] = (offset >> 24) & 0xFF
self._datacur += 4
def f_dbstr_print(dst, *args):
"""Print multiple string / number to dst.
:param dst: Destination address (Not EPD player)
:param args: Things to print
"""
if ut.isUnproxyInstance(dst, DBString):
dst = dst.GetStringMemoryAddr()
args = ut.FlattenList(args)
for arg in args:
if ut.isUnproxyInstance(arg, bytes):
dst = f_dbstr_addstr(dst, c.Db(arg + b'\0'))
elif ut.isUnproxyInstance(arg, str):
dst = f_dbstr_addstr(dst, c.Db(ut.u2b(arg) + b'\0'))
elif ut.isUnproxyInstance(arg, DBString):
dst = f_dbstr_addstr(dst, arg.GetStringMemoryAddr())
elif ut.isUnproxyInstance(arg, int):
# int and c.EUDVariable should act the same if possible.
# EUDVariable has a value of 32bit unsigned integer.
# So we adjust arg to be in the same range.
dst = f_dbstr_addstr(dst, c.Db(
ut.u2b(str(arg & 0xFFFFFFFF)) + b'\0'))
elif ut.isUnproxyInstance(arg, c.EUDVariable):
dst = f_dbstr_adddw(dst, arg)
elif c.IsConstExpr(arg):
dst = f_dbstr_adddw(dst, arg)
elif ut.isUnproxyInstance(arg, hptr):
dst = f_dbstr_addptr(dst, arg._value)
else:
raise ut.EPError(
'Object wit unknown parameter type %s given to f_eudprint.'
% type(arg)
)
return dst
def WritePayload(self, pbuffer):
"""Write object"""
raise ut.EPError('Override')
def GetDataSize(self):
"""Memory size of object."""
raise ut.EPError('Override')
def _GetLastLoopBlock():
for block in reversed(ut.EUDGetBlockList()):
if _IsLoopBlock(block[1]):
return block
raise ut.EPError('No loop block surrounding this code area')
def getValueAddr(self):
raise ut.EPError('override')
