diff --git a/src/ai/script_ai.cpp b/src/ai/script_ai.cpp index 7ddf2fcbf..f35130cae 100644 --- a/src/ai/script_ai.cpp +++ b/src/ai/script_ai.cpp @@ -1452,6 +1452,40 @@ static int CclAiProcessorSetup(lua_State *l) return 1; } +static CTCPSocket * AiProcessorSendState(lua_State *l, char prefix) +{ + LuaCheckArgs(l, 3); + CTCPSocket *s = (CTCPSocket *)lua_touserdata(l, 1); + if (s == NULL) { + LuaError(l, "first argument must be valid handle returned from a previous AiProcessorSetup call"); + } + + uint32_t reward = htonl(LuaToNumber(l, 2)); + if (!lua_istable(l, 3)) { + LuaError(l, "3rd argument to AiProcessorStep must be table"); + } + + char stepBuf[1029] = {'\0'}; // room for prefix + uint32 reward + 256 uint32 variables + stepBuf[0] = prefix; + int i = 1; + + memcpy(stepBuf + i, &reward, sizeof(uint32_t)); + i += sizeof(uint32_t); + + for (lua_pushnil(l); lua_next(l, 3); lua_pop(l, 1)) { + // idx is ignored + uint32_t var = htonl(LuaToNumber(l, -1)); + memcpy(stepBuf + i, &var, sizeof(uint32_t)); + i += sizeof(uint32_t); + if (i + sizeof(uint32_t) > 1025) { + LuaError(l, "too many state variables"); + } + } + s->Send(stepBuf, i); + + return s; +} + /** * AiProcessorStep(handle, reward_since_last_call, table_of_state_variables) */ @@ -1466,62 +1500,16 @@ static int CclAiProcessorStep(lua_State *l) // The next call to this function will be the updated state, reward for the // last action - LuaCheckArgs(l, 3); - CTCPSocket *s = (CTCPSocket *)lua_touserdata(l, 1); - if (s == NULL) { - LuaError(l, "first argument must be valid handle returned from a previous AiProcessorSetup call"); - } - - uint32_t reward = htonl(LuaToNumber(l, 2)); - char buf[5]; - buf[0] = 'R'; - memcpy(buf, &reward, sizeof(uint32_t)); - s->Send(buf, 5); - - if (!lua_istable(l, 3)) { - LuaError(l, "3rd argument to AiProcessorStep must be table"); - } - - char stepBuf[1025] = {'\0'}; // room for 256 variables - stepBuf[0] = 'S'; - int i = 1; - for (lua_pushnil(l); lua_next(l, 3); lua_pop(l, 1)) { - // idx is ignored - uint32_t var = htonl(LuaToNumber(l, -1)); - memcpy(stepBuf + i, &var, sizeof(uint32_t)); - i += sizeof(uint32_t); - if (i + sizeof(uint32_t) > 1025) { - LuaError(l, "too many state variables"); - } - } - s->Send(stepBuf, i); - + CTCPSocket *s = AiProcessorSendState(l, 'S'); int action = 0; s->Recv(&action, 1); lua_pushnumber(l, action); return 1; } -static int CclAiProcessorClose(lua_State *l) +static int CclAiProcessorEnd(lua_State *l) { - LuaCheckArgs(l, 2); - CTCPSocket *s = (CTCPSocket *)lua_touserdata(l, 1); - if (s == NULL) { - LuaError(l, "first argument must be valid handle returned from a previous AiProcessorSetup call"); - } - - int gameresult = LuaToNumber(l, 2); - switch (gameresult) { - case GameVictory: - s->Send("E\2", 2); - break; - case GameDefeat: - s->Send("E\1", 2); - break; - default: - s->Send("E\0", 2); - break; - } + CTCPSocket *s = AiProcessorSendState(l, 'E'); s->Close(); delete s; return 0; @@ -1574,7 +1562,7 @@ void AiCclRegister() // for external AI processors lua_register(Lua, "AiProcessorSetup", CclAiProcessorSetup); lua_register(Lua, "AiProcessorStep", CclAiProcessorStep); - lua_register(Lua, "AiProcessorClose", CclAiProcessorClose); + lua_register(Lua, "AiProcessorEnd", CclAiProcessorEnd); } //@}