<commit_before>/*
 * imconn.cpp
 *
 *  Created on: 2013-6-5
 *      Author: ziteng@mogujie.com
 */

#include "imconn.h"

//static uint64_t g_send_pkt_cnt = 0;		// 发送数据包总数
//static uint64_t g_recv_pkt_cnt = 0;		// 接收数据包总数

static CImConn* FindImConn(ConnMap_t* imconn_map, net_handle_t handle)
{
	CImConn* pConn = NULL;
	ConnMap_t::iterator iter = imconn_map->find(handle);
	if (iter != imconn_map->end())
	{
		pConn = iter->second;
		pConn->AddRef();
	}

	return pConn;
}

void imconn_callback(void* callback_data, uint8_t msg, uint32_t handle, void* pParam)
{
	NOTUSED_ARG(handle);
	NOTUSED_ARG(pParam);

	if (!callback_data)
		return;

	ConnMap_t* conn_map = (ConnMap_t*)callback_data;
	CImConn* pConn = FindImConn(conn_map, handle);
	if (!pConn)
		return;

	//log("msg=%d, handle=%d ", msg, handle);

	switch (msg)
	{
	case NETLIB_MSG_CONFIRM:
		pConn->OnConfirm();
		break;
	case NETLIB_MSG_READ:
		pConn->OnRead();
		break;
	case NETLIB_MSG_WRITE:
		pConn->OnWrite();
		break;
	case NETLIB_MSG_CLOSE:
		pConn->OnClose();
		break;
	case NETLIB_MSG_CLOSE_UDP:
		pConn->OnClose(); // UDP只有最后退出的时候close, 除非手动实现udp_listen
		break;
	case NETLIB_MSG_READ_UDP:
		pConn->OnReadUDP();
		break;
	case NETLIB_MSG_WRITE_UDP:
		pConn->OnWriteUDP();
		break;
	default:
		log("!!!imconn_callback error msg: %d ", msg);
		break;
	}

	pConn->ReleaseRef();
}

//////////////////////////
CImConn::CImConn()
{
	//log("CImConn::CImConn ");

	m_busy = false;
	m_handle = NETLIB_INVALID_HANDLE;
	m_recv_bytes = 0;

	m_last_send_tick = m_last_recv_tick = get_tick_count();
}

CImConn::~CImConn()
{
	//log("CImConn::~CImConn, handle=%d ", m_handle);
}

int CImConn::Send(void* data, int len)
{
	m_last_send_tick = get_tick_count();
//	++g_send_pkt_cnt;

	if (m_busy)
	{
		m_out_buf.Write(data, len);
		return len;
	}

	int offset = 0;
	int remain = len;
	while (remain > 0) {
		int send_size = remain;
		if (send_size > NETLIB_MAX_SOCKET_BUF_SIZE) {
			send_size = NETLIB_MAX_SOCKET_BUF_SIZE;
		}

		int ret = netlib_send(m_handle, (char*)data + offset , send_size);
		if (ret <= 0) {
			ret = 0;
			break;
		}

		offset += ret;
		remain -= ret;
	}

	if (remain > 0)
	{
		m_out_buf.Write((char*)data + offset, remain);
		m_busy = true;
		log("send busy, remain=%d ", m_out_buf.GetWriteOffset());
	}
    else
    {
        OnWriteCompelete();
    }

	return len;
}

void CImConn::OnRead()
{
	for (;;)
	{
		uint32_t free_buf_len = m_in_buf.GetAllocSize() - m_in_buf.GetWriteOffset();
		if (free_buf_len < READ_BUF_SIZE)
			m_in_buf.Extend(READ_BUF_SIZE);

		int ret = netlib_recv(m_handle, m_in_buf.GetBuffer() + m_in_buf.GetWriteOffset(), READ_BUF_SIZE);
		if (ret <= 0)
			break;

		m_recv_bytes += ret;
		m_in_buf.IncWriteOffset(ret);

		m_last_recv_tick = get_tick_count();
	}

    CImPdu* pPdu = NULL;
	try
    {
		while ( ( pPdu = CImPdu::ReadPdu(m_in_buf.GetBuffer(), m_in_buf.GetWriteOffset()) ) )
		{
            uint32_t pdu_len = pPdu->GetLength();

			HandlePdu(pPdu);

			m_in_buf.Read(NULL, pdu_len);
			delete pPdu;
            pPdu = NULL;
//			++g_recv_pkt_cnt;
		}
	} catch (CPduException& ex) {
		log("!!!catch exception, sid=%u, cid=%u, err_code=%u, err_msg=%s, close the connection ",
				ex.GetServiceId(), ex.GetCommandId(), ex.GetErrorCode(), ex.GetErrorMsg());
        if (pPdu) {
            delete pPdu;
            pPdu = NULL;
        }
        OnClose();
	}
}

void CImConn::OnWrite()
{
	if (!m_busy)
		return;

	while (m_out_buf.GetWriteOffset() > 0) {
		int send_size = m_out_buf.GetWriteOffset();
		if (send_size > NETLIB_MAX_SOCKET_BUF_SIZE) {
			send_size = NETLIB_MAX_SOCKET_BUF_SIZE;
		}

		int ret = netlib_send(m_handle, m_out_buf.GetBuffer(), send_size);
		if (ret <= 0) {
			ret = 0;
			break;
		}

		m_out_buf.Read(NULL, ret);
	}

	if (m_out_buf.GetWriteOffset() == 0) {
		m_busy = false;
	}

	log("onWrite, remain=%d ", m_out_buf.GetWriteOffset());
}

void CImConn::OnReadUDP()
{
	sockaddr_in sender; // 发送端的地址 从哪发来的
	char recvbuf[128]={0};
	memset(recvbuf,0,128);
	printf("CImConn OnReadUDP\n");
	socklen_t dwSender = sizeof(sender);
	int ret = recvfrom(m_sock_handle, (char *)&recvbuf, 128, 0, (sockaddr *)&sender, &dwSender);

	if(ret <= 0)
	{	// 读取要参照 void CImConn::OnRead()
		printf("recv error");
		return;
	}

	CImPdu* pPdu = NULL;
	try
    {
		while ( ( pPdu = CImPdu::ReadPdu(recvbuf, ret) ) )
		{
            uint32_t pdu_len = pPdu->GetLength();
			HandlePdu(pPdu);
			delete pPdu;
            pPdu = NULL;
		}
	} catch (CPduException& ex) {
		log("!!!catch exception, sid=%u, cid=%u, err_code=%u, err_msg=%s, close the connection ",
				ex.GetServiceId(), ex.GetCommandId(), ex.GetErrorCode(), ex.GetErrorMsg());
        if (pPdu) {
            delete pPdu;
            pPdu = NULL;
        }
        OnClose();
	}
}

void CImConn::OnWriteUDP()
{
}

<commit_msg>xie uchar_t<commit_after>/*
 * imconn.cpp
 *
 *  Created on: 2013-6-5
 *      Author: ziteng@mogujie.com
 */

#include "imconn.h"

//static uint64_t g_send_pkt_cnt = 0;		// 发送数据包总数
//static uint64_t g_recv_pkt_cnt = 0;		// 接收数据包总数

static CImConn* FindImConn(ConnMap_t* imconn_map, net_handle_t handle)
{
	CImConn* pConn = NULL;
	ConnMap_t::iterator iter = imconn_map->find(handle);
	if (iter != imconn_map->end())
	{
		pConn = iter->second;
		pConn->AddRef();
	}

	return pConn;
}

void imconn_callback(void* callback_data, uint8_t msg, uint32_t handle, void* pParam)
{
	NOTUSED_ARG(handle);
	NOTUSED_ARG(pParam);

	if (!callback_data)
		return;

	ConnMap_t* conn_map = (ConnMap_t*)callback_data;
	CImConn* pConn = FindImConn(conn_map, handle);
	if (!pConn)
		return;

	//log("msg=%d, handle=%d ", msg, handle);

	switch (msg)
	{
	case NETLIB_MSG_CONFIRM:
		pConn->OnConfirm();
		break;
	case NETLIB_MSG_READ:
		pConn->OnRead();
		break;
	case NETLIB_MSG_WRITE:
		pConn->OnWrite();
		break;
	case NETLIB_MSG_CLOSE:
		pConn->OnClose();
		break;
	case NETLIB_MSG_CLOSE_UDP:
		pConn->OnClose(); // UDP只有最后退出的时候close, 除非手动实现udp_listen
		break;
	case NETLIB_MSG_READ_UDP:
		pConn->OnReadUDP();
		break;
	case NETLIB_MSG_WRITE_UDP:
		pConn->OnWriteUDP();
		break;
	default:
		log("!!!imconn_callback error msg: %d ", msg);
		break;
	}

	pConn->ReleaseRef();
}

//////////////////////////
CImConn::CImConn()
{
	//log("CImConn::CImConn ");

	m_busy = false;
	m_handle = NETLIB_INVALID_HANDLE;
	m_recv_bytes = 0;

	m_last_send_tick = m_last_recv_tick = get_tick_count();
}

CImConn::~CImConn()
{
	//log("CImConn::~CImConn, handle=%d ", m_handle);
}

int CImConn::Send(void* data, int len)
{
	m_last_send_tick = get_tick_count();
//	++g_send_pkt_cnt;

	if (m_busy)
	{
		m_out_buf.Write(data, len);
		return len;
	}

	int offset = 0;
	int remain = len;
	while (remain > 0) {
		int send_size = remain;
		if (send_size > NETLIB_MAX_SOCKET_BUF_SIZE) {
			send_size = NETLIB_MAX_SOCKET_BUF_SIZE;
		}

		int ret = netlib_send(m_handle, (char*)data + offset , send_size);
		if (ret <= 0) {
			ret = 0;
			break;
		}

		offset += ret;
		remain -= ret;
	}

	if (remain > 0)
	{
		m_out_buf.Write((char*)data + offset, remain);
		m_busy = true;
		log("send busy, remain=%d ", m_out_buf.GetWriteOffset());
	}
    else
    {
        OnWriteCompelete();
    }

	return len;
}

void CImConn::OnRead()
{
	for (;;)
	{
		uint32_t free_buf_len = m_in_buf.GetAllocSize() - m_in_buf.GetWriteOffset();
		if (free_buf_len < READ_BUF_SIZE)
			m_in_buf.Extend(READ_BUF_SIZE);

		int ret = netlib_recv(m_handle, m_in_buf.GetBuffer() + m_in_buf.GetWriteOffset(), READ_BUF_SIZE);
		if (ret <= 0)
			break;

		m_recv_bytes += ret;
		m_in_buf.IncWriteOffset(ret);

		m_last_recv_tick = get_tick_count();
	}

    CImPdu* pPdu = NULL;
	try
    {
		while ( ( pPdu = CImPdu::ReadPdu(m_in_buf.GetBuffer(), m_in_buf.GetWriteOffset()) ) )
		{
            uint32_t pdu_len = pPdu->GetLength();

			HandlePdu(pPdu);

			m_in_buf.Read(NULL, pdu_len);
			delete pPdu;
            pPdu = NULL;
//			++g_recv_pkt_cnt;
		}
	} catch (CPduException& ex) {
		log("!!!catch exception, sid=%u, cid=%u, err_code=%u, err_msg=%s, close the connection ",
				ex.GetServiceId(), ex.GetCommandId(), ex.GetErrorCode(), ex.GetErrorMsg());
        if (pPdu) {
            delete pPdu;
            pPdu = NULL;
        }
        OnClose();
	}
}

void CImConn::OnWrite()
{
	if (!m_busy)
		return;

	while (m_out_buf.GetWriteOffset() > 0) {
		int send_size = m_out_buf.GetWriteOffset();
		if (send_size > NETLIB_MAX_SOCKET_BUF_SIZE) {
			send_size = NETLIB_MAX_SOCKET_BUF_SIZE;
		}

		int ret = netlib_send(m_handle, m_out_buf.GetBuffer(), send_size);
		if (ret <= 0) {
			ret = 0;
			break;
		}

		m_out_buf.Read(NULL, ret);
	}

	if (m_out_buf.GetWriteOffset() == 0) {
		m_busy = false;
	}

	log("onWrite, remain=%d ", m_out_buf.GetWriteOffset());
}

void CImConn::OnReadUDP()
{
	sockaddr_in sender; // 发送端的地址 从哪发来的
	uchar_t recvbuf[128]={0};
	memset(recvbuf,0,128);
	printf("CImConn OnReadUDP\n");
	socklen_t dwSender = sizeof(sender);
	int ret = recvfrom(m_sock_handle, (uchar_t *)&recvbuf, 128, 0, (sockaddr *)&sender, &dwSender);

	if(ret <= 0)
	{	// 读取要参照 void CImConn::OnRead()
		printf("recv error");
		return;
	}

	CImPdu* pPdu = NULL;
	try
    {
		while ( ( pPdu = CImPdu::ReadPdu(recvbuf, ret) ) )
		{
            uint32_t pdu_len = pPdu->GetLength();
			HandlePdu(pPdu);
			delete pPdu;
            pPdu = NULL;
		}
	} catch (CPduException& ex) {
		log("!!!catch exception, sid=%u, cid=%u, err_code=%u, err_msg=%s, close the connection ",
				ex.GetServiceId(), ex.GetCommandId(), ex.GetErrorCode(), ex.GetErrorMsg());
        if (pPdu) {
            delete pPdu;
            pPdu = NULL;
        }
        OnClose();
	}
}

void CImConn::OnWriteUDP()
{
}

<|endoftext|>


<commit_before>#include "Answer.h"

#include <string.h>
#include <stdlib.h>

namespace mDNSResolver {
  Answer::Answer() {
    this->name = NULL;
    this->data = NULL;
  }

  Answer::~Answer() {
    if(this->name) {
      free(this->name);
    }
    if(this->data) {
      free(this->data);
    }
  }

  MDNS_RESULT Answer::parseAnswer(unsigned char* buffer, unsigned int len, unsigned int* offset, Answer* answer) {
    char* assembled = (char *)malloc(sizeof(char) * MDNS_MAX_NAME_LEN);
    int nameLen = Answer::assembleName(buffer, len, offset, &assembled);

    if(nameLen == -1 * E_MDNS_POINTER_OVERFLOW) {
      free(assembled);
      return nameLen;
    }

    answer->name = (char *)malloc(sizeof(char) * nameLen);
    parseName(&answer->name, assembled, strlen(assembled));
    free(assembled);

    answer->type = (buffer[(*offset)++] << 8) + buffer[(*offset)++];

    unsigned int aclass = (buffer[(*offset)++] << 8) + buffer[(*offset)++];
    answer->cacheflush = aclass >> 15;
    answer->aclass = aclass & 0x7f;

    answer->ttl = (buffer[(*offset)++] << 24) + (buffer[(*offset)++] << 16) + (buffer[(*offset)++] << 8) + buffer[(*offset)++];
    answer->len = (buffer[(*offset)++] << 8) + buffer[(*offset)++];

    if(answer->type == MDNS_A_RECORD) {
      answer->data = (unsigned char *)malloc(sizeof(unsigned char) * answer->len);
      memcpy(answer->data, buffer + (*offset), answer->len);
      (*offset) += answer->len;
    } else if(answer->type == MDNS_CNAME_RECORD) {
      unsigned int dataOffset = (*offset);
      //(*offset) += answer->len;

      //char* assembled = (char *)malloc(sizeof(char) * MDNS_MAX_NAME_LEN);
      //MDNS_RESULT result = Answer::assembleName(buffer, len, &dataOffset, &assembled, answer->len);
      //if(result != E_MDNS_OK) {
        //free(assembled);
        //return result;
      //}
      //printf("Assembled: %s\n", assembled);
      //answer->len = parseName(answer->data, assembled, strlen(assembled));
      //free(assembled);
      answer->data = (unsigned char *)malloc(sizeof(unsigned char) * answer->len);
      memcpy(answer->data, buffer + (*offset), answer->len);
      (*offset) += answer->len;
    } else {
      // Not an A record or a CNAME. Ignore.
    }
    return E_MDNS_OK;
  }

  // Response should have the name set already
  MDNS_RESULT Answer::buildResponse(unsigned char* buffer, unsigned int len, Response& response) {
    if((buffer[2] & 0b10000000) != 0b10000000) {
      // Not an answer packet
      return E_MDNS_OK;
    }

    if(buffer[2] & 0b00000010) {
      // Truncated - we don't know what to do with these
      return E_MDNS_TRUNCATED;
    }

    if (buffer[3] & 0b00001111) {
      return E_MDNS_PACKET_ERROR;
    }

    unsigned int answerCount = (buffer[6] << 8) + buffer[7];

    // For this library, we are only interested in packets that contain answers
    if(answerCount == 0) {
      return E_MDNS_OK;
    }

    unsigned int offset = 0;

    MDNS_RESULT questionResult = skipQuestions(buffer, len, &offset);
    if(questionResult != E_MDNS_OK) {
      return questionResult;
    }

    Answer* answers = new Answer[answerCount];
    for(int i = 0; i < answerCount; i++) {
      parseAnswer(buffer, len, &offset, answers + i);
    }

    delete[] answers;
    return E_MDNS_OK;
  }

  // Converts a encoded DNS name into a FQDN.
  // name: pointer to char array where the result will be stored. Needs to have already been allocated. It's allocated length should be len - 1
  // mapped: The encoded DNS name
  // len: Length of mapped
  MDNS_RESULT Answer::parseName(char** name, const char* mapped, unsigned int len) {
    unsigned int namePointer = 0;
    unsigned int mapPointer = 0;

    while(mapPointer < len) {
      int labelLength = mapped[mapPointer++];

      if(namePointer + labelLength > len - 1) {
        return E_MDNS_INVALID_LABEL_LENGTH;
      }

      if(namePointer != 0) {
        (*name)[namePointer++] = '.';
      }

      for(int i = 0; i < labelLength; i++) {
        (*name)[namePointer++] = mapped[mapPointer++];
      }
    }

    (*name)[len - 1] = '\0';

    return E_MDNS_OK;
  }

  int Answer::assembleName(unsigned char *buffer, unsigned int len, unsigned int *offset, char **name, unsigned int maxlen) {
    unsigned int index = 0;
    unsigned int nameLength = 0;

    while(index < maxlen) {
      if((buffer[*offset] & 0xc0) == 0xc0) {
        unsigned int pointerOffset = ((buffer[(*offset)++] & 0x3f) << 8) + buffer[(*offset)++];
        if(pointerOffset > len) {
          // Points to somewhere beyond the packet
          return -1 * E_MDNS_POINTER_OVERFLOW;
        }

        char *namePointer = *name + nameLength;
        int pointerLen = assembleName(buffer, len, &pointerOffset, &namePointer, maxlen - nameLength);

        if(pointerLen < 0) {
          return pointerLen;
        }

        nameLength += pointerLen;

        break;
      } else if(buffer[*offset] == '\0') {
        (*name)[nameLength++] = buffer[(*offset)++];
        break;
      } else {
        (*name)[nameLength++] = buffer[(*offset)++];
      }
      index++;
    }

    return nameLength;
  }

  int Answer::assembleName(unsigned char *buffer, unsigned int len, unsigned int *offset, char **name) {
    return assembleName(buffer, len, offset, name, MDNS_MAX_NAME_LEN);
  }

  // Work out how many bytes are dedicated to questions. Since we aren't answering questions, they can be skipped
  // buffer: The mDNS packet we are parsing
  // len: Length of the packet
  // offset: the byte we are up to in the parsing process
  MDNS_RESULT Answer::skipQuestions(unsigned char* buffer, unsigned int len, unsigned int* offset) {
    unsigned int questionCount = (buffer[4] << 8) + buffer[5];

    *offset += 12;
    for(int i = 0; i < questionCount; i++) {

      while(buffer[*offset] != '\0') {
        // If it's a pointer, add two to the counter
        if((buffer[*offset] & 0xc0) == 0xc0) {
          (*offset) += 2;
          break;
        } else {
          unsigned int labelSize = (unsigned int)buffer[*offset];

          if(labelSize > 0x3f) {
            return E_MDNS_PACKET_ERROR;
          }

          (*offset) += 1; // Increment to move to the next byte
          (*offset) += labelSize;

          if(*offset > len) {
            return E_MDNS_PACKET_ERROR;
          }
        }
      }

      (*offset) += 5; // 2 bytes for the qtypes and 2 bytes qclass + plus one to land us on the next bit
    }

    if(*offset > len + 1) {
      return E_MDNS_PACKET_ERROR;
    }

    return E_MDNS_OK;
  }
};
<commit_msg>using MDNS_MAX_NAME_LEN when assembling name pointers<commit_after>#include "Answer.h"

#include <string.h>
#include <stdlib.h>

namespace mDNSResolver {
  Answer::Answer() {
    this->name = NULL;
    this->data = NULL;
  }

  Answer::~Answer() {
    if(this->name) {
      free(this->name);
    }
    if(this->data) {
      free(this->data);
    }
  }

  MDNS_RESULT Answer::parseAnswer(unsigned char* buffer, unsigned int len, unsigned int* offset, Answer* answer) {
    char* assembled = (char *)malloc(sizeof(char) * MDNS_MAX_NAME_LEN);
    int nameLen = Answer::assembleName(buffer, len, offset, &assembled);

    if(nameLen == -1 * E_MDNS_POINTER_OVERFLOW) {
      free(assembled);
      return nameLen;
    }

    answer->name = (char *)malloc(sizeof(char) * nameLen);
    parseName(&answer->name, assembled, strlen(assembled));
    free(assembled);

    answer->type = (buffer[(*offset)++] << 8) + buffer[(*offset)++];

    unsigned int aclass = (buffer[(*offset)++] << 8) + buffer[(*offset)++];
    answer->cacheflush = aclass >> 15;
    answer->aclass = aclass & 0x7f;

    answer->ttl = (buffer[(*offset)++] << 24) + (buffer[(*offset)++] << 16) + (buffer[(*offset)++] << 8) + buffer[(*offset)++];
    answer->len = (buffer[(*offset)++] << 8) + buffer[(*offset)++];

    if(answer->type == MDNS_A_RECORD) {
      answer->data = (unsigned char *)malloc(sizeof(unsigned char) * answer->len);
      memcpy(answer->data, buffer + (*offset), answer->len);
      (*offset) += answer->len;
    } else if(answer->type == MDNS_CNAME_RECORD) {
      unsigned int dataOffset = (*offset);
      //(*offset) += answer->len;

      //char* assembled = (char *)malloc(sizeof(char) * MDNS_MAX_NAME_LEN);
      //MDNS_RESULT result = Answer::assembleName(buffer, len, &dataOffset, &assembled, answer->len);
      //if(result != E_MDNS_OK) {
        //free(assembled);
        //return result;
      //}
      //printf("Assembled: %s\n", assembled);
      //answer->len = parseName(answer->data, assembled, strlen(assembled));
      //free(assembled);
      answer->data = (unsigned char *)malloc(sizeof(unsigned char) * answer->len);
      memcpy(answer->data, buffer + (*offset), answer->len);
      (*offset) += answer->len;
    } else {
      // Not an A record or a CNAME. Ignore.
    }
    return E_MDNS_OK;
  }

  // Response should have the name set already
  MDNS_RESULT Answer::buildResponse(unsigned char* buffer, unsigned int len, Response& response) {
    if((buffer[2] & 0b10000000) != 0b10000000) {
      // Not an answer packet
      return E_MDNS_OK;
    }

    if(buffer[2] & 0b00000010) {
      // Truncated - we don't know what to do with these
      return E_MDNS_TRUNCATED;
    }

    if (buffer[3] & 0b00001111) {
      return E_MDNS_PACKET_ERROR;
    }

    unsigned int answerCount = (buffer[6] << 8) + buffer[7];

    // For this library, we are only interested in packets that contain answers
    if(answerCount == 0) {
      return E_MDNS_OK;
    }

    unsigned int offset = 0;

    MDNS_RESULT questionResult = skipQuestions(buffer, len, &offset);
    if(questionResult != E_MDNS_OK) {
      return questionResult;
    }

    Answer* answers = new Answer[answerCount];
    for(int i = 0; i < answerCount; i++) {
      parseAnswer(buffer, len, &offset, answers + i);
    }

    delete[] answers;
    return E_MDNS_OK;
  }

  // Converts a encoded DNS name into a FQDN.
  // name: pointer to char array where the result will be stored. Needs to have already been allocated. It's allocated length should be len - 1
  // mapped: The encoded DNS name
  // len: Length of mapped
  MDNS_RESULT Answer::parseName(char** name, const char* mapped, unsigned int len) {
    unsigned int namePointer = 0;
    unsigned int mapPointer = 0;

    while(mapPointer < len) {
      int labelLength = mapped[mapPointer++];

      if(namePointer + labelLength > len - 1) {
        return E_MDNS_INVALID_LABEL_LENGTH;
      }

      if(namePointer != 0) {
        (*name)[namePointer++] = '.';
      }

      for(int i = 0; i < labelLength; i++) {
        (*name)[namePointer++] = mapped[mapPointer++];
      }
    }

    (*name)[len - 1] = '\0';

    return E_MDNS_OK;
  }

  int Answer::assembleName(unsigned char *buffer, unsigned int len, unsigned int *offset, char **name, unsigned int maxlen) {
    unsigned int index = 0;
    unsigned int nameLength = 0;

    while(index < maxlen) {
      if((buffer[*offset] & 0xc0) == 0xc0) {
        unsigned int pointerOffset = ((buffer[(*offset)++] & 0x3f) << 8) + buffer[(*offset)++];
        if(pointerOffset > len) {
          // Points to somewhere beyond the packet
          return -1 * E_MDNS_POINTER_OVERFLOW;
        }

        char *namePointer = *name + nameLength;
        int pointerLen = assembleName(buffer, len, &pointerOffset, &namePointer, MDNS_MAX_NAME_LEN - nameLength);

        if(pointerLen < 0) {
          return pointerLen;
        }

        nameLength += pointerLen;

        break;
      } else if(buffer[*offset] == '\0') {
        (*name)[nameLength++] = buffer[(*offset)++];
        break;
      } else {
        (*name)[nameLength++] = buffer[(*offset)++];
      }
      index++;
    }

    return nameLength;
  }

  int Answer::assembleName(unsigned char *buffer, unsigned int len, unsigned int *offset, char **name) {
    return assembleName(buffer, len, offset, name, MDNS_MAX_NAME_LEN);
  }

  // Work out how many bytes are dedicated to questions. Since we aren't answering questions, they can be skipped
  // buffer: The mDNS packet we are parsing
  // len: Length of the packet
  // offset: the byte we are up to in the parsing process
  MDNS_RESULT Answer::skipQuestions(unsigned char* buffer, unsigned int len, unsigned int* offset) {
    unsigned int questionCount = (buffer[4] << 8) + buffer[5];

    *offset += 12;
    for(int i = 0; i < questionCount; i++) {

      while(buffer[*offset] != '\0') {
        // If it's a pointer, add two to the counter
        if((buffer[*offset] & 0xc0) == 0xc0) {
          (*offset) += 2;
          break;
        } else {
          unsigned int labelSize = (unsigned int)buffer[*offset];

          if(labelSize > 0x3f) {
            return E_MDNS_PACKET_ERROR;
          }

          (*offset) += 1; // Increment to move to the next byte
          (*offset) += labelSize;

          if(*offset > len) {
            return E_MDNS_PACKET_ERROR;
          }
        }
      }

      (*offset) += 5; // 2 bytes for the qtypes and 2 bytes qclass + plus one to land us on the next bit
    }

    if(*offset > len + 1) {
      return E_MDNS_PACKET_ERROR;
    }

    return E_MDNS_OK;
  }
};
<|endoftext|>


<commit_before>/* ************************************************
 * GEANT4 VCGLIB/CAD INTERFACE
 *
 * File:      CADMesh.cc
 *
 * Author:    Christopher M Poole,
 * Email:     mail@christopherpoole.net
 *
 * Date:      7th March, 2011
 **************************************************/

// USER //
#include "CADMesh.hh"

// GEANT4 //
#include "G4String.hh"
#include "G4UIcommand.hh"
#include "G4ThreeVector.hh"
#include "G4LogicalVolume.hh"
#include "G4TessellatedSolid.hh"
#include "G4TriangularFacet.hh"
#include "G4Tet.hh"
#include "G4ThreeVector.hh"

#ifndef NOVCGLIB
// VCGLIB //
#include "wrap/io_trimesh/import.h"
#include "wrap/io_trimesh/import_ply.h"
#include "wrap/io_trimesh/import_stl.h"
#include "wrap/io_trimesh/import_off.h"

#include "wrap/ply/plylib.h"
#endif

#ifndef NOVCGLIB
using namespace vcg::tri::io;
#endif

CADMesh::CADMesh(char * file_name, char * file_type, double units, G4ThreeVector offset, G4bool reverse)
{
    units_ = units;
    offset_ = offset;
    reverse_ = reverse;

    file_name_ = file_name;
    file_type_ = file_type;
    file_type_.toUpper();

    material_ = NULL;
    quality_ = 0;

    has_mesh_ = false;
    has_solid_ = false;
    verbose_ = 0;
}

CADMesh::CADMesh(char * file_name, char * file_type, G4Material * material, double quality)
{
    units_ = mm;
    offset_ = G4ThreeVector();
    reverse_ = false;

    file_name_ = file_name;
    file_type_ = file_type;
    file_type_.toUpper();

    material_ = material;
    quality_ = quality;

    has_mesh_ = false;
    has_solid_ = false;
    verbose_ = 0;
}

CADMesh::CADMesh(char * file_name, char * file_type)
{
    units_ = mm;
    offset_ = G4ThreeVector();
    reverse_ = false;

    file_name_ = file_name;
    file_type_ = file_type;
    file_type_.toUpper();

    material_ = NULL;
    quality_ = 0;

    has_mesh_ = false;
    has_solid_ = false;
    verbose_ = 0;
}

CADMesh::CADMesh(char * file_name, char * file_type, G4Material * material)
{
    units_ = mm;
    offset_ = G4ThreeVector();
    reverse_ = false;

    file_name_ = file_name;
    file_type_ = file_type;
    file_type_.toUpper();

    material_ = material;
    quality_ = 0;

    has_mesh_ = false;
    has_solid_ = false;
    verbose_ = 0;
}

CADMesh::CADMesh(char * file_name, char * file_type, G4Material * material, double quality, G4ThreeVector offset)
{
    units_ = mm;
    offset_ = offset;
    reverse_ = false;

    file_name_ = file_name;
    file_type_ = file_type;
    file_type_.toUpper();

    material_ = material;
    quality_ = quality;

    has_mesh_ = false;
    has_solid_ = false;
    verbose_ = 0;
}

CADMesh::~CADMesh()
{
}

#ifndef NOVCGLIB
G4VSolid* CADMesh::TessellatedMesh()
{
    if (!has_mesh_) {
        if (file_type_ == "STL") {
            ImporterSTL<CADTriMesh>::Open(m, file_name_);
        } else if (file_type_ == "PLY") {
            ImporterPLY<CADTriMesh>::Open(m, file_name_);
        } else if (file_type_ == "OFF") {
            ImporterOFF<CADTriMesh>::Open(m, file_name_);
        } else {
            G4cerr << "CADMesh/LoadSTL: "
                   << "No G4TessellatedSoild to return. Specify valid mesh type (STL, PLY, OFF), not: "
                   << file_type_
                   << G4endl;
            has_mesh_ = false;

            return 0;
        }
        has_mesh_ = true;

    } else {
        G4cerr << "CADMesh/LoadSTL: "
               << "Mesh already loaded from "
               << file_name_
               << ", not loading. Use CADMesh/GetSolid to get the currently loaded mesh as a G4TessellatedSolid"
               << G4endl;
        return 0;
    }

    if (!has_mesh_) {
        G4cerr << "CADMesh/BuildSolid: "
               << "Load a mesh of type STL, PLY or OFF first."
               << G4endl;
        return 0;
    }

    volume_solid = new G4TessellatedSolid(file_name_);

    G4ThreeVector point_1 = G4ThreeVector();
    G4ThreeVector point_2 = G4ThreeVector();
    G4ThreeVector point_3 = G4ThreeVector();

    CADTriMesh::FaceIterator face_iterator;

    for(face_iterator=m.face.begin(); face_iterator!=m.face.end(); ++face_iterator)
    {
        point_1 = G4ThreeVector(
                (*face_iterator).V(0)->P()[0] * units_ - offset_.x(),
                (*face_iterator).V(0)->P()[1] * units_ - offset_.y(),
                (*face_iterator).V(0)->P()[2] * units_ - offset_.z());
        point_2 = G4ThreeVector(
                (*face_iterator).V(1)->P()[0] * units_ - offset_.x(),
                (*face_iterator).V(1)->P()[1] * units_ - offset_.y(),
                (*face_iterator).V(1)->P()[2] * units_ - offset_.z());
        point_3 = G4ThreeVector(
                (*face_iterator).V(2)->P()[0] * units_ - offset_.x(),
                (*face_iterator).V(2)->P()[1] * units_ - offset_.y(),
                (*face_iterator).V(2)->P()[2] * units_ - offset_.z());

        G4TriangularFacet * facet;
        if (reverse_ == false) {
            facet = new G4TriangularFacet(point_1, point_2, point_3, ABSOLUTE);
        } else {
            facet = new G4TriangularFacet(point_2, point_1, point_3, ABSOLUTE);
        }
        volume_solid->AddFacet((G4VFacet*) facet);
    }

    volume_solid->SetSolidClosed(true);

    return volume_solid;
}
#endif

#ifndef NOTET
G4AssemblyVolume * CADMesh::TetrahedralMesh()
{
    // USAGE: assembly->MakeImprint(world_logical, assembly_transform_3d, 0); //

    G4bool do_tet = true;
    if (file_type_ == "STL") {
        G4bool state = in.load_stl(file_name_);
    } else if (file_type_ == "PLY") {
        G4bool state = in.load_ply(file_name_);
    } else if (file_type_ == "TET") {
        G4bool state = out.load_tetmesh(file_name_);
        do_tet = false;
    } else if (file_type_ == "OFF") {
        G4bool state = out.load_off(file_name_);
        do_tet = false;
    }

    if (do_tet)
    {
        G4String config = G4String("Yp");
        if (quality_ > 0) config = config + G4String("q") + G4UIcommand::ConvertToString(quality_);
#ifdef DEBUG
        G4cout << "Tetrahedralisation configuration: " << config << G4endl;
#endif
        tetrahedralize((char *) config.c_str(), &in, &out);
    }

#ifdef DEBUG
    G4cout << "Tetrahedra available: " << out.numberoftetrahedra << G4endl;
#endif

    assembly = new G4AssemblyVolume();
    G4RotationMatrix * element_rotation = new G4RotationMatrix();
    G4ThreeVector element_position = G4ThreeVector();
    G4Transform3D assembly_transform = G4Translate3D();

    for (int i=0; i<out.numberoftetrahedra; i++) {
        int index_offset = i * 4; /* For a tetrahedron, out.numberofcorners == 4 */

        G4ThreeVector p1 = GetTetPoint(index_offset);
        G4ThreeVector p2 = GetTetPoint(index_offset + 1);
        G4ThreeVector p3 = GetTetPoint(index_offset + 2);
        G4ThreeVector p4 = GetTetPoint(index_offset + 3);

        G4String tet_name = file_name_ + G4String("_tet_") + G4UIcommand::ConvertToString(i);

        G4VSolid * tet_solid = new G4Tet(tet_name + G4String("_solid"), p1, p2, p3, p4, 0);
        G4LogicalVolume * tet_logical = new G4LogicalVolume(tet_solid, material_, tet_name + G4String("_logical"), 0, 0, 0);
        assembly->AddPlacedVolume(tet_logical, element_position, element_rotation);

#ifdef DEBUG
        if (i%1000 == 0) G4cout << "Tetrahedrons added: " << i << G4endl;
#endif
    }

#ifdef DEBUG
        G4cout << "Loading of " << out.numberoftetrahedra << "tetrahedra complete." << G4endl;
#endif

    return assembly;
}

G4ThreeVector CADMesh::GetTetPoint(G4int index_offset)
{
    return G4ThreeVector(out.pointlist[out.tetrahedronlist[index_offset]*3] - offset_.x(),
                         out.pointlist[out.tetrahedronlist[index_offset]*3+1] - offset_.y(),
                         out.pointlist[out.tetrahedronlist[index_offset]*3+2] - offset_.z());
}
#endif
<commit_msg>prints material name for tetmesh<commit_after>/* ************************************************
 * GEANT4 VCGLIB/CAD INTERFACE
 *
 * File:      CADMesh.cc
 *
 * Author:    Christopher M Poole,
 * Email:     mail@christopherpoole.net
 *
 * Date:      7th March, 2011
 **************************************************/

// USER //
#include "CADMesh.hh"

// GEANT4 //
#include "G4String.hh"
#include "G4UIcommand.hh"
#include "G4ThreeVector.hh"
#include "G4LogicalVolume.hh"
#include "G4TessellatedSolid.hh"
#include "G4TriangularFacet.hh"
#include "G4Tet.hh"
#include "G4ThreeVector.hh"
#include "G4Material.hh"

#ifndef NOVCGLIB
// VCGLIB //
#include "wrap/io_trimesh/import.h"
#include "wrap/io_trimesh/import_ply.h"
#include "wrap/io_trimesh/import_stl.h"
#include "wrap/io_trimesh/import_off.h"

#include "wrap/ply/plylib.h"
#endif

#ifndef NOVCGLIB
using namespace vcg::tri::io;
#endif

CADMesh::CADMesh(char * file_name, char * file_type, double units, G4ThreeVector offset, G4bool reverse)
{
    units_ = units;
    offset_ = offset;
    reverse_ = reverse;

    file_name_ = file_name;
    file_type_ = file_type;
    file_type_.toUpper();

    material_ = NULL;
    quality_ = 0;

    has_mesh_ = false;
    has_solid_ = false;
    verbose_ = 0;
}

CADMesh::CADMesh(char * file_name, char * file_type, G4Material * material, double quality)
{
    units_ = mm;
    offset_ = G4ThreeVector();
    reverse_ = false;

    file_name_ = file_name;
    file_type_ = file_type;
    file_type_.toUpper();

    material_ = material;
    quality_ = quality;

    has_mesh_ = false;
    has_solid_ = false;
    verbose_ = 0;
}

CADMesh::CADMesh(char * file_name, char * file_type)
{
    units_ = mm;
    offset_ = G4ThreeVector();
    reverse_ = false;

    file_name_ = file_name;
    file_type_ = file_type;
    file_type_.toUpper();

    material_ = NULL;
    quality_ = 0;

    has_mesh_ = false;
    has_solid_ = false;
    verbose_ = 0;
}

CADMesh::CADMesh(char * file_name, char * file_type, G4Material * material)
{
    units_ = mm;
    offset_ = G4ThreeVector();
    reverse_ = false;

    file_name_ = file_name;
    file_type_ = file_type;
    file_type_.toUpper();

    material_ = material;
    quality_ = 0;

    has_mesh_ = false;
    has_solid_ = false;
    verbose_ = 0;
}

CADMesh::CADMesh(char * file_name, char * file_type, G4Material * material, double quality, G4ThreeVector offset)
{
    units_ = mm;
    offset_ = offset;
    reverse_ = false;

    file_name_ = file_name;
    file_type_ = file_type;
    file_type_.toUpper();

    material_ = material;
    quality_ = quality;

    has_mesh_ = false;
    has_solid_ = false;
    verbose_ = 0;
}

CADMesh::~CADMesh()
{
}

#ifndef NOVCGLIB
G4VSolid* CADMesh::TessellatedMesh()
{
    if (!has_mesh_) {
        if (file_type_ == "STL") {
            ImporterSTL<CADTriMesh>::Open(m, file_name_);
        } else if (file_type_ == "PLY") {
            ImporterPLY<CADTriMesh>::Open(m, file_name_);
        } else if (file_type_ == "OFF") {
            ImporterOFF<CADTriMesh>::Open(m, file_name_);
        } else {
            G4cerr << "CADMesh/LoadSTL: "
                   << "No G4TessellatedSoild to return. Specify valid mesh type (STL, PLY, OFF), not: "
                   << file_type_
                   << G4endl;
            has_mesh_ = false;

            return 0;
        }
        has_mesh_ = true;

    } else {
        G4cerr << "CADMesh/LoadSTL: "
               << "Mesh already loaded from "
               << file_name_
               << ", not loading. Use CADMesh/GetSolid to get the currently loaded mesh as a G4TessellatedSolid"
               << G4endl;
        return 0;
    }

    if (!has_mesh_) {
        G4cerr << "CADMesh/BuildSolid: "
               << "Load a mesh of type STL, PLY or OFF first."
               << G4endl;
        return 0;
    }

    volume_solid = new G4TessellatedSolid(file_name_);

    G4ThreeVector point_1 = G4ThreeVector();
    G4ThreeVector point_2 = G4ThreeVector();
    G4ThreeVector point_3 = G4ThreeVector();

    CADTriMesh::FaceIterator face_iterator;

    for(face_iterator=m.face.begin(); face_iterator!=m.face.end(); ++face_iterator)
    {
        point_1 = G4ThreeVector(
                (*face_iterator).V(0)->P()[0] * units_ - offset_.x(),
                (*face_iterator).V(0)->P()[1] * units_ - offset_.y(),
                (*face_iterator).V(0)->P()[2] * units_ - offset_.z());
        point_2 = G4ThreeVector(
                (*face_iterator).V(1)->P()[0] * units_ - offset_.x(),
                (*face_iterator).V(1)->P()[1] * units_ - offset_.y(),
                (*face_iterator).V(1)->P()[2] * units_ - offset_.z());
        point_3 = G4ThreeVector(
                (*face_iterator).V(2)->P()[0] * units_ - offset_.x(),
                (*face_iterator).V(2)->P()[1] * units_ - offset_.y(),
                (*face_iterator).V(2)->P()[2] * units_ - offset_.z());

        G4TriangularFacet * facet;
        if (reverse_ == false) {
            facet = new G4TriangularFacet(point_1, point_2, point_3, ABSOLUTE);
        } else {
            facet = new G4TriangularFacet(point_2, point_1, point_3, ABSOLUTE);
        }
        volume_solid->AddFacet((G4VFacet*) facet);
    }

    volume_solid->SetSolidClosed(true);

    return volume_solid;
}
#endif

#ifndef NOTET
G4AssemblyVolume * CADMesh::TetrahedralMesh()
{
    // USAGE: assembly->MakeImprint(world_logical, assembly_transform_3d, 0); //

    G4bool do_tet = true;
    if (file_type_ == "STL") {
        G4bool state = in.load_stl(file_name_);
    } else if (file_type_ == "PLY") {
        G4bool state = in.load_ply(file_name_);
    } else if (file_type_ == "TET") {
        G4bool state = out.load_tetmesh(file_name_);
        do_tet = false;
    } else if (file_type_ == "OFF") {
        G4bool state = out.load_off(file_name_);
        do_tet = false;
    }

    if (do_tet)
    {
        G4String config = G4String("Yp");
        if (quality_ > 0) config = config + G4String("q") + G4UIcommand::ConvertToString(quality_);
#ifdef DEBUG
        G4cout << "Tetrahedralisation configuration: " << config << G4endl;
#endif
        tetrahedralize((char *) config.c_str(), &in, &out);
    }

#ifdef DEBUG
    G4cout << "Tetrahedra available: " << out.numberoftetrahedra << G4endl;
#endif

    assembly = new G4AssemblyVolume();
    G4RotationMatrix * element_rotation = new G4RotationMatrix();
    G4ThreeVector element_position = G4ThreeVector();
    G4Transform3D assembly_transform = G4Translate3D();

    for (int i=0; i<out.numberoftetrahedra; i++) {
        int index_offset = i * 4; /* For a tetrahedron, out.numberofcorners == 4 */

        G4ThreeVector p1 = GetTetPoint(index_offset);
        G4ThreeVector p2 = GetTetPoint(index_offset + 1);
        G4ThreeVector p3 = GetTetPoint(index_offset + 2);
        G4ThreeVector p4 = GetTetPoint(index_offset + 3);

        G4String tet_name = file_name_ + G4String("_tet_") + G4UIcommand::ConvertToString(i);

        G4VSolid * tet_solid = new G4Tet(tet_name + G4String("_solid"), p1, p2, p3, p4, 0);
        G4LogicalVolume * tet_logical = new G4LogicalVolume(tet_solid, material_, tet_name + G4String("_logical"), 0, 0, 0);
        assembly->AddPlacedVolume(tet_logical, element_position, element_rotation);

#ifdef DEBUG
        if (i%1000 == 0) G4cout << "Tetrahedrons added: " << i << G4endl;
#endif
    }

#ifdef DEBUG
    G4cout << "Loading of " << out.numberoftetrahedra << " tetrahedra complete with material: " << material_->GetName() << G4endl;
#endif

    return assembly;
}

G4ThreeVector CADMesh::GetTetPoint(G4int index_offset)
{
    return G4ThreeVector(out.pointlist[out.tetrahedronlist[index_offset]*3] - offset_.x(),
                         out.pointlist[out.tetrahedronlist[index_offset]*3+1] - offset_.y(),
                         out.pointlist[out.tetrahedronlist[index_offset]*3+2] - offset_.z());
}
#endif
<|endoftext|>


<commit_before>/*
 *   CrissCross
 *   A multi-purpose cross-platform library.
 *
 *   A product of Uplink Laboratories.
 *
 *   (c) 2006-2007 Steven Noonan.
 *   Licensed under the New BSD License.
 *
 */

#include <crisscross/universal_include.h>
#include <crisscross/debug.h>
#include <crisscross/mutex.h>

namespace CrissCross
{
	namespace System
	{
		Mutex::Mutex ()
		{
			m_lockCount = 0;
		#ifdef TARGET_OS_WINDOWS
			InitializeCriticalSection(&m_criticalSection);
		#else
			pthread_mutexattr_init ( &m_mutexAttr );
			pthread_mutexattr_settype ( &m_mutexAttr, PTHREAD_MUTEX_ERRORCHECK );
			int error = pthread_mutex_init ( &m_hMutex, &m_mutexAttr );
			CoreAssert ( error == 0 );
		#endif
		}

		Mutex::~Mutex ()
		{
			CoreAssert ( m_lockCount == 0 );
			Unlock();
		#ifdef TARGET_OS_WINDOWS
			DeleteCriticalSection(&m_criticalSection);
		#else
			pthread_mutex_destroy(&m_hMutex);
		#endif
		}

		void
		Mutex::Lock ()
		{
			m_lockCount++;
		#ifdef TARGET_OS_WINDOWS
			EnterCriticalSection(&m_criticalSection);
		#else
			int error = pthread_mutex_lock(&m_hMutex);
			CoreAssert ( error == 0 );
		#endif
		}

		void
		Mutex::Unlock ()
		{
			m_lockCount--;
		#ifdef TARGET_OS_WINDOWS
			LeaveCriticalSection(&m_criticalSection);
		#else
			int error = pthread_mutex_unlock(&m_hMutex);
			CoreAssert ( error == 0 );
		#endif
		}
	}
}
<commit_msg>Recursive mutex enabled.<commit_after>/*
 *   CrissCross
 *   A multi-purpose cross-platform library.
 *
 *   A product of Uplink Laboratories.
 *
 *   (c) 2006-2007 Steven Noonan.
 *   Licensed under the New BSD License.
 *
 */

#include <crisscross/universal_include.h>
#include <crisscross/debug.h>
#include <crisscross/mutex.h>

namespace CrissCross
{
	namespace System
	{
		Mutex::Mutex ()
		{
			m_lockCount = 0;
		#ifdef TARGET_OS_WINDOWS
			InitializeCriticalSection(&m_criticalSection);
		#else
			pthread_mutexattr_init ( &m_mutexAttr );
			pthread_mutexattr_settype ( &m_mutexAttr, PTHREAD_MUTEX_RECURSIVE );
			int error = pthread_mutex_init ( &m_hMutex, &m_mutexAttr );
			CoreAssert ( error == 0 );
		#endif
		}

		Mutex::~Mutex ()
		{
			CoreAssert ( m_lockCount == 0 );
			Unlock();
		#ifdef TARGET_OS_WINDOWS
			DeleteCriticalSection(&m_criticalSection);
		#else
			pthread_mutex_destroy(&m_hMutex);
		#endif
		}

		void
		Mutex::Lock ()
		{
			m_lockCount++;
		#ifdef TARGET_OS_WINDOWS
			EnterCriticalSection(&m_criticalSection);
		#else
			int error = pthread_mutex_lock(&m_hMutex);
			CoreAssert ( error == 0 );
		#endif
		}

		void
		Mutex::Unlock ()
		{
			m_lockCount--;
		#ifdef TARGET_OS_WINDOWS
			LeaveCriticalSection(&m_criticalSection);
		#else
			int error = pthread_mutex_unlock(&m_hMutex);
			CoreAssert ( error == 0 );
		#endif
		}
	}
}
<|endoftext|>


<commit_before>#include "McMaker/TpcEffFitter.h"
#include "Common.h"

// ROOT
#include "TGraphAsymmErrors.h"
#include "TFile.h"
#include "TLatex.h"

#include "Logger.h"
#include "Reporter.h"
#include "RooPlotLib.h"
using namespace jdb;

void TpcEffFitter::initialize(  ){

	DEBUG( classname(), "( " << config.getFilename() << ", " << nodePath << " )" )
	outputPath = config.getString( nodePath + ".output:path", "" );

	book = unique_ptr<HistoBook>( new HistoBook( outputPath + config.getString( nodePath +  ".output.data", "TpcEff.root" ), config, "", "" ) );
}

void TpcEffFitter::make(){
	DEBUG(classname(), "")

	RooPlotLib rpl;

	gStyle->SetOptFit( 111 );
	string params_file =  config.getString( nodePath + ".output.params" );
	if ( "" == params_file ){
		ERROR( classname(), "Specifiy an output params file for the parameters" )
		return;
	}

	ofstream out( params_file.c_str() );

	out << "<?xml version=\"1.0\" encoding=\"UTF-8\"?>" << endl;
	out << "<config>" << endl;

	vector<string> labels = config.getStringVector( nodePath + ".CentralityLabels" );
	vector< int> cbins = config.getIntVector( nodePath + ".CentralityBins" );
	Reporter rp( config, nodePath + ".Reporter." );

	double minP0Error = config.getDouble( nodePath + ".Systematics:minP0Error" );
	TLatex Tl; Tl.SetTextFont(43); Tl.SetTextSize(20);

	DEBUG( "Starting plc loop" )
	for ( string plc : Common::species ){
		if ( "E" == plc || "D" == plc )	// skip Electrons and Deuterons if I've got those in the list
			continue;

		for ( string c : Common::sCharges ){

			out << "\t<" << plc << "_" << c << ">" << endl;

			string fnMc = config.getString( nodePath + ".input:url" ) + "TpcEff_" + plc + "_" + c + "_mc" + ".root";
			TFile * fmc = new TFile( fnMc.c_str(), "READ" );
			string fnRc = config.getString( nodePath + ".input:url" ) + "TpcEff_" + plc + "_" + c + "_rc" + ".root";
			TFile * frc = new TFile( fnRc.c_str(), "READ" );

			DEBUG( classname(), "Mc File = " << fmc );
			DEBUG( classname(), "Rc File = " << frc );

			if ( !fmc->IsOpen() || !frc->IsOpen() )
				continue;

			// build an efficiency for each centrality
			for ( int b : cbins ){

				TH1 * hMc = (TH1*)fmc->Get( ("inclusive/pt_" + ts( b ) + "_" + c ).c_str() );
				TH1 * hRc = (TH1*)frc->Get( ("inclusive/pt_" + ts( b ) + "_" + c ).c_str() );

				INFO( classname(), "N bins MC = " << hMc->GetNbinsX() );
				INFO( classname(), "N bins RC = " << hRc->GetNbinsX() );

				for ( int i = 0; i <= hMc->GetNbinsX() + 1; i++ ){
					if ( hMc->GetBinContent( i ) < hRc->GetBinContent( i ) ){
						// set to 100%
						if ( i > 5 )
							hMc->SetBinContent( i, hRc->GetBinContent( i ) );
						else{
							hRc->SetBinContent( i, 0 );
							hMc->SetBinContent( i, 0 );
						}
					}
				}

				hMc->Sumw2();
				hRc->Sumw2();

				TGraphAsymmErrors g;

				g.SetName( (plc + "_" + c + "_" + ts(b)).c_str() );
				g.BayesDivide( hRc, hMc );

				INFO( classname(), "Adding TGraphAsymmErrors to book" );
				book->add( plc + "_" + c + "_" + ts(b),  &g );

				// do the fit
				TF1 * fitFunc = new TF1( "effFitFunc", "[0] * exp( - pow( [1] / x, [2] ) )", 0.0, 5.0 );
				fitFunc->SetParameters( .85, 0.05, 5.0, -0.05 );
				fitFunc->SetParLimits( 0, 0.5, 1.0 );
				fitFunc->SetParLimits( 1, 0.0, 0.5 );
				fitFunc->SetParLimits( 2, 0.0, 100000 );

				// fist fit shape
				TFitResultPtr fitPointer = g.Fit( fitFunc, "RSWW" );
				fitPointer = g.Fit( fitFunc, "RS" );

				// ensure that uncertainty on efficiency is at least X%
				if ( fitFunc->GetParError( 0 ) < minP0Error ) {
					INFO( classname(), "P0 uncertainty below threshold (" << (minP0Error * 100) << "%" );
					INFO( classname(), "Setting P0 error to " << (minP0Error * 100) << "%" );
					fitFunc->SetParError( 0, minP0Error );
				}

				INFO( classname(), "FitPointer = " << fitPointer );
				TGraphErrors * band = Common::choleskyBands( fitPointer, fitFunc, 5000, 200, &rp );

				rp.newPage();
				rpl.style( &g ).set( "title", Common::plc_label( plc, c ) + " : " + labels[ b ] + ", Fit 68%CL(Red Band)" )
					.set( "yr", 0, 1.1 ).set( "optfit", 111 )
					.set( "xr", 0, 4.5 )
					.set("y", "Efficiency x Acceptance")
					.set( "x", "p_{T}^{MC} [GeV/c]" )
					.draw();

				INFO( classname(), "Stat Box" );
				gStyle->SetStatY( 0.5 );
				gStyle->SetStatX( 0.85 );

				fitFunc->SetLineColor( kRed );
				fitFunc->Draw("same");

				INFO( classname(), "Drawing CL band" );

				band->SetFillColorAlpha( kRed, 0.7 );
				band->Draw( "same e3" );

				Tl.DrawLatex( 0.5, 0.5, "[0] * #exp( - pow( [1] / x, [2] ) )" );

				rp.savePage();
				string imgName = outputPath + "/img/TpcEff_" + plc + "_" + c + "_" + ts(b) + ".png";
				INFO( classname(), "Exporting image to : " << imgName );
				rp.saveImage( imgName );

				INFO( classname(), "Exporting Params" );
				exportParams( b, fitFunc, fitPointer, out );

			} // loop centrality bins

			out << "\t</" << plc << "_" << c << ">" << endl;

		} // loop on charge
	} // loop on plc

	out << "</config>" << endl;
	out.close();

}

void TpcEffFitter::exportParams( int bin, TF1 * f, TFitResultPtr result, ofstream &out ){
	INFO( classname(), "(bin=" << bin << ", f=" << f << ", fPtr=" << result << " )" )
	out << "\t\t<TpcEffParams bin=\"" << bin << "\" ";
	out << Common::toXml( f, result );
	out << "/>" << endl;
}<commit_msg>Adding eq to printout<commit_after>#include "McMaker/TpcEffFitter.h"
#include "Common.h"

// ROOT
#include "TGraphAsymmErrors.h"
#include "TFile.h"
#include "TLatex.h"

#include "Logger.h"
#include "Reporter.h"
#include "RooPlotLib.h"
using namespace jdb;

void TpcEffFitter::initialize(  ){

	DEBUG( classname(), "( " << config.getFilename() << ", " << nodePath << " )" )
	outputPath = config.getString( nodePath + ".output:path", "" );

	book = unique_ptr<HistoBook>( new HistoBook( outputPath + config.getString( nodePath +  ".output.data", "TpcEff.root" ), config, "", "" ) );
}

void TpcEffFitter::make(){
	DEBUG(classname(), "")

	RooPlotLib rpl;

	gStyle->SetOptFit( 111 );
	string params_file =  config.getString( nodePath + ".output.params" );
	if ( "" == params_file ){
		ERROR( classname(), "Specifiy an output params file for the parameters" )
		return;
	}

	ofstream out( params_file.c_str() );

	out << "<?xml version=\"1.0\" encoding=\"UTF-8\"?>" << endl;
	out << "<config>" << endl;

	vector<string> labels = config.getStringVector( nodePath + ".CentralityLabels" );
	vector< int> cbins = config.getIntVector( nodePath + ".CentralityBins" );
	Reporter rp( config, nodePath + ".Reporter." );

	double minP0Error = config.getDouble( nodePath + ".Systematics:minP0Error" );
	TLatex Tl; Tl.SetTextFont(43); Tl.SetTextSize(20);

	DEBUG( "Starting plc loop" )
	for ( string plc : Common::species ){
		if ( "E" == plc || "D" == plc )	// skip Electrons and Deuterons if I've got those in the list
			continue;

		for ( string c : Common::sCharges ){

			out << "\t<" << plc << "_" << c << ">" << endl;

			string fnMc = config.getString( nodePath + ".input:url" ) + "TpcEff_" + plc + "_" + c + "_mc" + ".root";
			TFile * fmc = new TFile( fnMc.c_str(), "READ" );
			string fnRc = config.getString( nodePath + ".input:url" ) + "TpcEff_" + plc + "_" + c + "_rc" + ".root";
			TFile * frc = new TFile( fnRc.c_str(), "READ" );

			DEBUG( classname(), "Mc File = " << fmc );
			DEBUG( classname(), "Rc File = " << frc );

			if ( !fmc->IsOpen() || !frc->IsOpen() )
				continue;

			// build an efficiency for each centrality
			for ( int b : cbins ){

				TH1 * hMc = (TH1*)fmc->Get( ("inclusive/pt_" + ts( b ) + "_" + c ).c_str() );
				TH1 * hRc = (TH1*)frc->Get( ("inclusive/pt_" + ts( b ) + "_" + c ).c_str() );

				INFO( classname(), "N bins MC = " << hMc->GetNbinsX() );
				INFO( classname(), "N bins RC = " << hRc->GetNbinsX() );

				for ( int i = 0; i <= hMc->GetNbinsX() + 1; i++ ){
					if ( hMc->GetBinContent( i ) < hRc->GetBinContent( i ) ){
						// set to 100%
						if ( i > 5 )
							hMc->SetBinContent( i, hRc->GetBinContent( i ) );
						else{
							hRc->SetBinContent( i, 0 );
							hMc->SetBinContent( i, 0 );
						}
					}
				}

				hMc->Sumw2();
				hRc->Sumw2();

				TGraphAsymmErrors g;

				g.SetName( (plc + "_" + c + "_" + ts(b)).c_str() );
				g.BayesDivide( hRc, hMc );

				INFO( classname(), "Adding TGraphAsymmErrors to book" );
				book->add( plc + "_" + c + "_" + ts(b),  &g );

				// do the fit
				TF1 * fitFunc = new TF1( "effFitFunc", "[0] * exp( - pow( [1] / x, [2] ) )", 0.0, 5.0 );
				fitFunc->SetParameters( .85, 0.05, 5.0, -0.05 );
				fitFunc->SetParLimits( 0, 0.5, 1.0 );
				fitFunc->SetParLimits( 1, 0.0, 0.5 );
				fitFunc->SetParLimits( 2, 0.0, 100000 );

				// fist fit shape
				TFitResultPtr fitPointer = g.Fit( fitFunc, "RSWW" );
				fitPointer = g.Fit( fitFunc, "RS" );

				// ensure that uncertainty on efficiency is at least X%
				if ( fitFunc->GetParError( 0 ) < minP0Error ) {
					INFO( classname(), "P0 uncertainty below threshold (" << (minP0Error * 100) << "%" );
					INFO( classname(), "Setting P0 error to " << (minP0Error * 100) << "%" );
					fitFunc->SetParError( 0, minP0Error );
				}

				INFO( classname(), "FitPointer = " << fitPointer );
				TGraphErrors * band = Common::choleskyBands( fitPointer, fitFunc, 5000, 200, &rp );

				rp.newPage();
				rpl.style( &g ).set( "title", Common::plc_label( plc, c ) + " : " + labels[ b ] + ", Fit 68%CL(Red Band)" )
					.set( "yr", 0, 1.1 ).set( "optfit", 111 )
					.set( "xr", 0, 4.5 )
					.set("y", "Efficiency x Acceptance")
					.set( "x", "p_{T}^{MC} [GeV/c]" )
					.draw();

				INFO( classname(), "Stat Box" );
				gStyle->SetStatY( 0.5 );
				gStyle->SetStatX( 0.85 );

				fitFunc->SetLineColor( kRed );
				fitFunc->Draw("same");

				INFO( classname(), "Drawing CL band" );

				band->SetFillColorAlpha( kRed, 0.7 );
				band->Draw( "same e3" );

				Tl.DrawLatex( 0.5, 0.5, "f(pT) = [p0]e^{ (#frac{-[p1]}{pT})^{[p2]} }" );

				rp.savePage();
				string imgName = outputPath + "/img/TpcEff_" + plc + "_" + c + "_" + ts(b) + ".png";
				INFO( classname(), "Exporting image to : " << imgName );
				rp.saveImage( imgName );

				INFO( classname(), "Exporting Params" );
				exportParams( b, fitFunc, fitPointer, out );

			} // loop centrality bins

			out << "\t</" << plc << "_" << c << ">" << endl;

		} // loop on charge
	} // loop on plc

	out << "</config>" << endl;
	out.close();

}

void TpcEffFitter::exportParams( int bin, TF1 * f, TFitResultPtr result, ofstream &out ){
	INFO( classname(), "(bin=" << bin << ", f=" << f << ", fPtr=" << result << " )" )
	out << "\t\t<TpcEffParams bin=\"" << bin << "\" ";
	out << Common::toXml( f, result );
	out << "/>" << endl;
}<|endoftext|>


<commit_before>// Copyright (c) 2010 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <deque>
#include <stdlib.h>

#include "gfx/rect.h"
#include "media/base/video_frame.h"
#include "remoting/base/codec_test.h"
#include "remoting/base/encoder.h"
#include "remoting/base/mock_objects.h"
#include "testing/gtest/include/gtest/gtest.h"

static const int kWidth = 320;
static const int kHeight = 240;
static const int kBytesPerPixel = 4;

// Some sample rects for testing.
static const gfx::Rect kTestRects[] = {
  gfx::Rect(0, 0, kWidth, kHeight),
  gfx::Rect(0, 0, kWidth / 2, kHeight / 2),
  gfx::Rect(kWidth / 2, kHeight / 2, kWidth / 2, kHeight / 2),
  gfx::Rect(16, 16, 16, 16),
  gfx::Rect(128, 64, 32, 32)
};

namespace remoting {

// A class to test that the state transition of the Encoder is correct.
class EncoderStateTester {
 public:
  EncoderStateTester()
      : next_state_(Encoder::EncodingStarting) {
  }

  ~EncoderStateTester() {
    EXPECT_EQ(Encoder::EncodingStarting, next_state_);
  }

  // Set the state output of the Encoder.
  void ReceivedState(Encoder::EncodingState state) {
    if (state & Encoder::EncodingStarting) {
      EXPECT_EQ(Encoder::EncodingStarting, next_state_);
      next_state_ = Encoder::EncodingInProgress | Encoder::EncodingEnded;
    } else {
      EXPECT_FALSE(next_state_ & Encoder::EncodingStarting);
    }

    if (state & Encoder::EncodingInProgress) {
      EXPECT_TRUE(next_state_ & Encoder::EncodingInProgress);
    }

    if (state & Encoder::EncodingEnded) {
      EXPECT_TRUE(next_state_ & Encoder::EncodingEnded);
      next_state_ = Encoder::EncodingStarting;
    }
  }

 private:
  Encoder::EncodingState next_state_;

  DISALLOW_COPY_AND_ASSIGN(EncoderStateTester);
};

// A class to test the message output of the encoder.
class EncoderMessageTester {
 public:
  EncoderMessageTester()
      : begin_rect_(0),
        rect_data_(0),
        end_rect_(0),
        state_(kWaitingForBeginRect),
        strict_(false) {
  }

  ~EncoderMessageTester() {
    EXPECT_EQ(begin_rect_, end_rect_);
    EXPECT_EQ(kWaitingForBeginRect, state_);
  }

  // Test that we received the correct message.
  void ReceivedMessage(HostMessage* message) {
    EXPECT_TRUE(message->has_update_stream_packet());

    if (state_ == kWaitingForBeginRect) {
      EXPECT_TRUE(message->update_stream_packet().has_begin_rect());
      state_ = kWaitingForRectData;
      ++begin_rect_;

      if (strict_) {
        gfx::Rect rect = rects_.front();
        rects_.pop_front();
        EXPECT_EQ(rect.x(), message->update_stream_packet().begin_rect().x());
        EXPECT_EQ(rect.y(), message->update_stream_packet().begin_rect().y());
        EXPECT_EQ(rect.width(),
                  message->update_stream_packet().begin_rect().width());
        EXPECT_EQ(rect.height(),
                  message->update_stream_packet().begin_rect().height());
      }
    } else {
      EXPECT_FALSE(message->update_stream_packet().has_begin_rect());
    }

    if (state_ == kWaitingForRectData) {
      if (message->update_stream_packet().has_rect_data()) {
        ++rect_data_;
      }

      if (message->update_stream_packet().has_end_rect()) {
        // Expect that we have received some data.
        EXPECT_GT(rect_data_, 0);
        rect_data_ = 0;
        state_ = kWaitingForBeginRect;
        ++end_rect_;
      }
    }
  }

  void set_strict(bool strict) {
    strict_ = strict;
  }

  void AddRects(const gfx::Rect* rects, int count) {
    rects_.insert(rects_.begin() + rects_.size(), rects, rects + count);
  }

 private:
  enum State {
    kWaitingForBeginRect,
    kWaitingForRectData,
  };

  int begin_rect_;
  int rect_data_;
  int end_rect_;
  State state_;
  bool strict_;

  std::deque<gfx::Rect> rects_;

  DISALLOW_COPY_AND_ASSIGN(EncoderMessageTester);
};

class DecoderTester {
 public:
  DecoderTester(Decoder* decoder)
      : strict_(false),
        decoder_(decoder),
        decode_done_(false) {
    media::VideoFrame::CreateFrame(media::VideoFrame::RGB32,
                                   kWidth, kHeight,
                                   base::TimeDelta(),
                                   base::TimeDelta(), &frame_);
    EXPECT_TRUE(frame_.get());
  }

  void ReceivedMessage(HostMessage* message) {
    if (message->has_update_stream_packet()) {
      decoder_->PartialDecode(message);
      return;
    }

    if (message->has_begin_update_stream()) {
      decoder_->BeginDecode(
          frame_, &update_rects_,
          NewRunnableMethod(this, &DecoderTester::OnPartialDecodeDone),
          NewRunnableMethod(this, &DecoderTester::OnDecodeDone));
    }

    if (message->has_end_update_stream()) {
      decoder_->EndDecode();
    }
    delete message;
    return;
  }

  void set_strict(bool strict) {
    strict_ = strict;
  }

  void set_capture_data(scoped_refptr<CaptureData> data) {
    capture_data_ = data;
  }

  void AddRects(const gfx::Rect* rects, int count) {
    rects_.insert(rects_.begin() + rects_.size(), rects, rects + count);
  }

  bool decode_done() const { return decode_done_; }
  void reset_decode_done() { decode_done_ = false; }

 private:
  void OnPartialDecodeDone() {
    if (!strict_)
      return;
    for (size_t i = 0; i < update_rects_.size(); ++i) {
      EXPECT_FALSE(rects_.empty());
      gfx::Rect rect = rects_.front();
      rects_.pop_front();
      EXPECT_EQ(rect, update_rects_[i]);
    }
  }

  void OnDecodeDone() {
    decode_done_ = true;
    if (!strict_)
      return;

    EXPECT_TRUE(capture_data_.get());
    for (int i = 0; i < DataPlanes::kPlaneCount; ++i) {
      if (!frame_->data(i) || !capture_data_->data_planes().data[i])
        continue;
      // TODO(hclam): HAndle YUV.
      int size = capture_data_->data_planes().strides[i] * kHeight;
      EXPECT_EQ(0, memcmp(capture_data_->data_planes().data[i],
                          frame_->data(i), size));
    }
  }

  bool strict_;
  std::deque<gfx::Rect> rects_;
  UpdatedRects update_rects_;
  Decoder* decoder_;
  scoped_refptr<media::VideoFrame> frame_;
  scoped_refptr<CaptureData> capture_data_;
  bool decode_done_;

  DISALLOW_COPY_AND_ASSIGN(DecoderTester);
};

class EncoderTester {
 public:
  EncoderTester(EncoderMessageTester* message_tester,
                EncoderStateTester* state_tester)
      : message_tester_(message_tester),
        state_tester_(state_tester),
        decoder_tester_(NULL),
        data_available_(0) {
  }

  ~EncoderTester() {
    EXPECT_GT(data_available_, 0);
  }

  void DataAvailable(HostMessage* message,
                     Encoder::EncodingState state) {
    ++data_available_;
    message_tester_->ReceivedMessage(message);
    state_tester_->ReceivedState(state);

    // Send the message to the DecoderTester.
    if (decoder_tester_) {
      if (state & Encoder::EncodingStarting) {
        HostMessage* begin_update = new HostMessage();
        begin_update->mutable_begin_update_stream();
        decoder_tester_->ReceivedMessage(begin_update);
      }

      if (state & Encoder::EncodingInProgress) {
        decoder_tester_->ReceivedMessage(message);
      }

      if (state & Encoder::EncodingEnded) {
        HostMessage* end_update = new HostMessage();
        end_update->mutable_end_update_stream();
        decoder_tester_->ReceivedMessage(end_update);
      }
	} else {
      delete message;
	}
  }

  void AddRects(const gfx::Rect* rects, int count) {
    message_tester_->AddRects(rects, count);
  }

  void set_decoder_tester(DecoderTester* decoder_tester) {
    decoder_tester_ = decoder_tester;
  }

 private:
  EncoderMessageTester* message_tester_;
  EncoderStateTester* state_tester_;
  DecoderTester* decoder_tester_;
  int data_available_;

  DISALLOW_COPY_AND_ASSIGN(EncoderTester);
};

scoped_refptr<CaptureData> PrepareEncodeData(PixelFormat format,
                                             uint8** memory) {
  // TODO(hclam): Support also YUV format.
  CHECK(format == PixelFormatRgb32);
  int size = kWidth * kHeight * kBytesPerPixel;

  *memory = new uint8[size];
  srand(0);
  for (int i = 0; i < size; ++i) {
    (*memory)[i] = rand() % 256;
  }

  DataPlanes planes;
  memset(planes.data, 0, sizeof(planes.data));
  memset(planes.strides, 0, sizeof(planes.strides));
  planes.data[0] = *memory;
  planes.strides[0] = kWidth * kBytesPerPixel;

  scoped_refptr<CaptureData> data =
      new CaptureData(planes, kWidth, kHeight, format);
  return data;
}

static void TestEncodingRects(Encoder* encoder,
                              EncoderTester* tester,
                              scoped_refptr<CaptureData> data,
                              const gfx::Rect* rects, int count) {
  data->mutable_dirty_rects().clear();
  data->mutable_dirty_rects().insert(
      data->mutable_dirty_rects().begin(), rects, rects + count);
  tester->AddRects(rects, count);

  encoder->Encode(data, true,
                  NewCallback(tester, &EncoderTester::DataAvailable));
}

void TestEncoder(Encoder* encoder, bool strict) {
  EncoderMessageTester message_tester;
  message_tester.set_strict(strict);

  EncoderStateTester state_tester;
  EncoderTester tester(&message_tester, &state_tester);

  uint8* memory;
  scoped_refptr<CaptureData> data =
      PrepareEncodeData(PixelFormatRgb32, &memory);

  TestEncodingRects(encoder, &tester, data, kTestRects, 1);
  TestEncodingRects(encoder, &tester, data, kTestRects + 1, 1);
  TestEncodingRects(encoder, &tester, data, kTestRects + 2, 1);
  TestEncodingRects(encoder, &tester, data, kTestRects + 3, 2);
  delete memory;
}

static void TestEncodingRects(Encoder* encoder,
                              EncoderTester* encoder_tester,
                              DecoderTester* decoder_tester,
                              scoped_refptr<CaptureData> data,
                              const gfx::Rect* rects, int count) {
  data->mutable_dirty_rects().clear();
  data->mutable_dirty_rects().insert(
      data->mutable_dirty_rects().begin(), rects, rects + count);
  encoder_tester->AddRects(rects, count);
  decoder_tester->AddRects(rects, count);
  decoder_tester->reset_decode_done();

  encoder->Encode(data, true,
                  NewCallback(encoder_tester, &EncoderTester::DataAvailable));
  EXPECT_TRUE(decoder_tester->decode_done());
}

void TestEncoderDecoder(Encoder* encoder, Decoder* decoder, bool strict) {
  EncoderMessageTester message_tester;
  message_tester.set_strict(strict);

  EncoderStateTester state_tester;
  EncoderTester encoder_tester(&message_tester, &state_tester);

  uint8* memory;
  scoped_refptr<CaptureData> data =
      PrepareEncodeData(PixelFormatRgb32, &memory);
  DecoderTester decoder_tester(decoder);
  decoder_tester.set_strict(strict);
  decoder_tester.set_capture_data(data);
  encoder_tester.set_decoder_tester(&decoder_tester);

  TestEncodingRects(encoder, &encoder_tester, &decoder_tester, data,
                    kTestRects, 1);
  TestEncodingRects(encoder, &encoder_tester, &decoder_tester, data,
                    kTestRects + 1, 1);
  TestEncodingRects(encoder, &encoder_tester, &decoder_tester, data,
                    kTestRects + 2, 1);
  TestEncodingRects(encoder, &encoder_tester, &decoder_tester, data,
                    kTestRects + 3, 2);
  delete memory;
}

}  // namespace remoting

DISABLE_RUNNABLE_METHOD_REFCOUNT(remoting::DecoderTester);
<commit_msg>Fix mismatched new[]/delete in EncoderVerbatimTest.TestEncoder + Fix presubmit warnings about tabs near line 270 + svn pset svn:eol-style LF remoting/base/codec_test.cc BUG=50126 TBR=hclam TEST=bots Review URL: http://codereview.chromium.org/3051012<commit_after>// Copyright (c) 2010 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <deque>
#include <stdlib.h>

#include "gfx/rect.h"
#include "media/base/video_frame.h"
#include "remoting/base/codec_test.h"
#include "remoting/base/encoder.h"
#include "remoting/base/mock_objects.h"
#include "testing/gtest/include/gtest/gtest.h"

static const int kWidth = 320;
static const int kHeight = 240;
static const int kBytesPerPixel = 4;

// Some sample rects for testing.
static const gfx::Rect kTestRects[] = {
  gfx::Rect(0, 0, kWidth, kHeight),
  gfx::Rect(0, 0, kWidth / 2, kHeight / 2),
  gfx::Rect(kWidth / 2, kHeight / 2, kWidth / 2, kHeight / 2),
  gfx::Rect(16, 16, 16, 16),
  gfx::Rect(128, 64, 32, 32)
};

namespace remoting {

// A class to test that the state transition of the Encoder is correct.
class EncoderStateTester {
 public:
  EncoderStateTester()
      : next_state_(Encoder::EncodingStarting) {
  }

  ~EncoderStateTester() {
    EXPECT_EQ(Encoder::EncodingStarting, next_state_);
  }

  // Set the state output of the Encoder.
  void ReceivedState(Encoder::EncodingState state) {
    if (state & Encoder::EncodingStarting) {
      EXPECT_EQ(Encoder::EncodingStarting, next_state_);
      next_state_ = Encoder::EncodingInProgress | Encoder::EncodingEnded;
    } else {
      EXPECT_FALSE(next_state_ & Encoder::EncodingStarting);
    }

    if (state & Encoder::EncodingInProgress) {
      EXPECT_TRUE(next_state_ & Encoder::EncodingInProgress);
    }

    if (state & Encoder::EncodingEnded) {
      EXPECT_TRUE(next_state_ & Encoder::EncodingEnded);
      next_state_ = Encoder::EncodingStarting;
    }
  }

 private:
  Encoder::EncodingState next_state_;

  DISALLOW_COPY_AND_ASSIGN(EncoderStateTester);
};

// A class to test the message output of the encoder.
class EncoderMessageTester {
 public:
  EncoderMessageTester()
      : begin_rect_(0),
        rect_data_(0),
        end_rect_(0),
        state_(kWaitingForBeginRect),
        strict_(false) {
  }

  ~EncoderMessageTester() {
    EXPECT_EQ(begin_rect_, end_rect_);
    EXPECT_EQ(kWaitingForBeginRect, state_);
  }

  // Test that we received the correct message.
  void ReceivedMessage(HostMessage* message) {
    EXPECT_TRUE(message->has_update_stream_packet());

    if (state_ == kWaitingForBeginRect) {
      EXPECT_TRUE(message->update_stream_packet().has_begin_rect());
      state_ = kWaitingForRectData;
      ++begin_rect_;

      if (strict_) {
        gfx::Rect rect = rects_.front();
        rects_.pop_front();
        EXPECT_EQ(rect.x(), message->update_stream_packet().begin_rect().x());
        EXPECT_EQ(rect.y(), message->update_stream_packet().begin_rect().y());
        EXPECT_EQ(rect.width(),
                  message->update_stream_packet().begin_rect().width());
        EXPECT_EQ(rect.height(),
                  message->update_stream_packet().begin_rect().height());
      }
    } else {
      EXPECT_FALSE(message->update_stream_packet().has_begin_rect());
    }

    if (state_ == kWaitingForRectData) {
      if (message->update_stream_packet().has_rect_data()) {
        ++rect_data_;
      }

      if (message->update_stream_packet().has_end_rect()) {
        // Expect that we have received some data.
        EXPECT_GT(rect_data_, 0);
        rect_data_ = 0;
        state_ = kWaitingForBeginRect;
        ++end_rect_;
      }
    }
  }

  void set_strict(bool strict) {
    strict_ = strict;
  }

  void AddRects(const gfx::Rect* rects, int count) {
    rects_.insert(rects_.begin() + rects_.size(), rects, rects + count);
  }

 private:
  enum State {
    kWaitingForBeginRect,
    kWaitingForRectData,
  };

  int begin_rect_;
  int rect_data_;
  int end_rect_;
  State state_;
  bool strict_;

  std::deque<gfx::Rect> rects_;

  DISALLOW_COPY_AND_ASSIGN(EncoderMessageTester);
};

class DecoderTester {
 public:
  DecoderTester(Decoder* decoder)
      : strict_(false),
        decoder_(decoder),
        decode_done_(false) {
    media::VideoFrame::CreateFrame(media::VideoFrame::RGB32,
                                   kWidth, kHeight,
                                   base::TimeDelta(),
                                   base::TimeDelta(), &frame_);
    EXPECT_TRUE(frame_.get());
  }

  void ReceivedMessage(HostMessage* message) {
    if (message->has_update_stream_packet()) {
      decoder_->PartialDecode(message);
      return;
    }

    if (message->has_begin_update_stream()) {
      decoder_->BeginDecode(
          frame_, &update_rects_,
          NewRunnableMethod(this, &DecoderTester::OnPartialDecodeDone),
          NewRunnableMethod(this, &DecoderTester::OnDecodeDone));
    }

    if (message->has_end_update_stream()) {
      decoder_->EndDecode();
    }
    delete message;
    return;
  }

  void set_strict(bool strict) {
    strict_ = strict;
  }

  void set_capture_data(scoped_refptr<CaptureData> data) {
    capture_data_ = data;
  }

  void AddRects(const gfx::Rect* rects, int count) {
    rects_.insert(rects_.begin() + rects_.size(), rects, rects + count);
  }

  bool decode_done() const { return decode_done_; }
  void reset_decode_done() { decode_done_ = false; }

 private:
  void OnPartialDecodeDone() {
    if (!strict_)
      return;
    for (size_t i = 0; i < update_rects_.size(); ++i) {
      EXPECT_FALSE(rects_.empty());
      gfx::Rect rect = rects_.front();
      rects_.pop_front();
      EXPECT_EQ(rect, update_rects_[i]);
    }
  }

  void OnDecodeDone() {
    decode_done_ = true;
    if (!strict_)
      return;

    EXPECT_TRUE(capture_data_.get());
    for (int i = 0; i < DataPlanes::kPlaneCount; ++i) {
      if (!frame_->data(i) || !capture_data_->data_planes().data[i])
        continue;
      // TODO(hclam): HAndle YUV.
      int size = capture_data_->data_planes().strides[i] * kHeight;
      EXPECT_EQ(0, memcmp(capture_data_->data_planes().data[i],
                          frame_->data(i), size));
    }
  }

  bool strict_;
  std::deque<gfx::Rect> rects_;
  UpdatedRects update_rects_;
  Decoder* decoder_;
  scoped_refptr<media::VideoFrame> frame_;
  scoped_refptr<CaptureData> capture_data_;
  bool decode_done_;

  DISALLOW_COPY_AND_ASSIGN(DecoderTester);
};

class EncoderTester {
 public:
  EncoderTester(EncoderMessageTester* message_tester,
                EncoderStateTester* state_tester)
      : message_tester_(message_tester),
        state_tester_(state_tester),
        decoder_tester_(NULL),
        data_available_(0) {
  }

  ~EncoderTester() {
    EXPECT_GT(data_available_, 0);
  }

  void DataAvailable(HostMessage* message,
                     Encoder::EncodingState state) {
    ++data_available_;
    message_tester_->ReceivedMessage(message);
    state_tester_->ReceivedState(state);

    // Send the message to the DecoderTester.
    if (decoder_tester_) {
      if (state & Encoder::EncodingStarting) {
        HostMessage* begin_update = new HostMessage();
        begin_update->mutable_begin_update_stream();
        decoder_tester_->ReceivedMessage(begin_update);
      }

      if (state & Encoder::EncodingInProgress) {
        decoder_tester_->ReceivedMessage(message);
      }

      if (state & Encoder::EncodingEnded) {
        HostMessage* end_update = new HostMessage();
        end_update->mutable_end_update_stream();
        decoder_tester_->ReceivedMessage(end_update);
      }
    } else {
      delete message;
    }
  }

  void AddRects(const gfx::Rect* rects, int count) {
    message_tester_->AddRects(rects, count);
  }

  void set_decoder_tester(DecoderTester* decoder_tester) {
    decoder_tester_ = decoder_tester;
  }

 private:
  EncoderMessageTester* message_tester_;
  EncoderStateTester* state_tester_;
  DecoderTester* decoder_tester_;
  int data_available_;

  DISALLOW_COPY_AND_ASSIGN(EncoderTester);
};

scoped_refptr<CaptureData> PrepareEncodeData(PixelFormat format,
                                             uint8** memory) {
  // TODO(hclam): Support also YUV format.
  CHECK(format == PixelFormatRgb32);
  int size = kWidth * kHeight * kBytesPerPixel;

  *memory = new uint8[size];
  srand(0);
  for (int i = 0; i < size; ++i) {
    (*memory)[i] = rand() % 256;
  }

  DataPlanes planes;
  memset(planes.data, 0, sizeof(planes.data));
  memset(planes.strides, 0, sizeof(planes.strides));
  planes.data[0] = *memory;
  planes.strides[0] = kWidth * kBytesPerPixel;

  scoped_refptr<CaptureData> data =
      new CaptureData(planes, kWidth, kHeight, format);
  return data;
}

static void TestEncodingRects(Encoder* encoder,
                              EncoderTester* tester,
                              scoped_refptr<CaptureData> data,
                              const gfx::Rect* rects, int count) {
  data->mutable_dirty_rects().clear();
  data->mutable_dirty_rects().insert(
      data->mutable_dirty_rects().begin(), rects, rects + count);
  tester->AddRects(rects, count);

  encoder->Encode(data, true,
                  NewCallback(tester, &EncoderTester::DataAvailable));
}

void TestEncoder(Encoder* encoder, bool strict) {
  EncoderMessageTester message_tester;
  message_tester.set_strict(strict);

  EncoderStateTester state_tester;
  EncoderTester tester(&message_tester, &state_tester);

  uint8* memory;
  scoped_refptr<CaptureData> data =
      PrepareEncodeData(PixelFormatRgb32, &memory);

  TestEncodingRects(encoder, &tester, data, kTestRects, 1);
  TestEncodingRects(encoder, &tester, data, kTestRects + 1, 1);
  TestEncodingRects(encoder, &tester, data, kTestRects + 2, 1);
  TestEncodingRects(encoder, &tester, data, kTestRects + 3, 2);
  delete [] memory;
}

static void TestEncodingRects(Encoder* encoder,
                              EncoderTester* encoder_tester,
                              DecoderTester* decoder_tester,
                              scoped_refptr<CaptureData> data,
                              const gfx::Rect* rects, int count) {
  data->mutable_dirty_rects().clear();
  data->mutable_dirty_rects().insert(
      data->mutable_dirty_rects().begin(), rects, rects + count);
  encoder_tester->AddRects(rects, count);
  decoder_tester->AddRects(rects, count);
  decoder_tester->reset_decode_done();

  encoder->Encode(data, true,
                  NewCallback(encoder_tester, &EncoderTester::DataAvailable));
  EXPECT_TRUE(decoder_tester->decode_done());
}

void TestEncoderDecoder(Encoder* encoder, Decoder* decoder, bool strict) {
  EncoderMessageTester message_tester;
  message_tester.set_strict(strict);

  EncoderStateTester state_tester;
  EncoderTester encoder_tester(&message_tester, &state_tester);

  uint8* memory;
  scoped_refptr<CaptureData> data =
      PrepareEncodeData(PixelFormatRgb32, &memory);
  DecoderTester decoder_tester(decoder);
  decoder_tester.set_strict(strict);
  decoder_tester.set_capture_data(data);
  encoder_tester.set_decoder_tester(&decoder_tester);

  TestEncodingRects(encoder, &encoder_tester, &decoder_tester, data,
                    kTestRects, 1);
  TestEncodingRects(encoder, &encoder_tester, &decoder_tester, data,
                    kTestRects + 1, 1);
  TestEncodingRects(encoder, &encoder_tester, &decoder_tester, data,
                    kTestRects + 2, 1);
  TestEncodingRects(encoder, &encoder_tester, &decoder_tester, data,
                    kTestRects + 3, 2);
  delete [] memory;
}

}  // namespace remoting

DISABLE_RUNNABLE_METHOD_REFCOUNT(remoting::DecoderTester);
<|endoftext|>


<commit_before>/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/stream_executor/tpu/tpu_executable_interface.h"

#include <utility>

#include "absl/algorithm/container.h"
#include "tensorflow/compiler/xla/service/maybe_owning_device_memory.h"
#include "tensorflow/compiler/xla/service/shaped_buffer.h"
#include "tensorflow/compiler/xla/service/transfer_manager.h"
#include "tensorflow/compiler/xla/shape_util.h"
#include "tensorflow/compiler/xla/status_macros.h"
#include "tensorflow/compiler/xla/util.h"

namespace xla {

namespace {

// Write the tuple index buffers (arrays of pointers).
static Status PopulateResultTupleBuffers(const ShapedBuffer& result,
                                         se::Stream* stream,
                                         se::Stream* transfer_stream) {
  TF_ASSIGN_OR_RETURN(auto transfer_manager, TransferManager::GetForPlatform(
                                                 stream->parent()->platform()));
  if (transfer_manager->CanShapedBufferBeAccessedNow(stream->parent(),
                                                     result)) {
    TF_RETURN_IF_ERROR(transfer_manager->WriteTupleIndexTablesAsync(
        transfer_stream ? transfer_stream : stream, result));
    if (transfer_stream && transfer_stream != stream) {
      stream->ThenWaitFor(transfer_stream);
    }
    return Status::OK();
  } else {
    return transfer_manager->WriteTupleIndexTablesAsync(stream, result);
  }
}

}  // namespace

StatusOr<ExecutionOutput>
TpuExecutableInterface::AllocateOutputMemoryWithInputReuse(
    const Shape& host_shape, const HloInputOutputAliasConfig& alias_config,
    se::DeviceMemoryAllocator* allocator,
    std::vector<ExecutionInput>* arguments, se::Stream* stream,
    se::Stream* transfer_stream) {
  auto stream_exec = stream->parent();
  auto device_ordinal = stream_exec->device_ordinal();
  VLOG(3) << "AllocateOutputMemoryWithInputReuse, device = " << device_ordinal
          << " host_shape = " << ShapeUtil::HumanStringWithLayout(host_shape);
  Shape device_shape = HostShapeToDeviceShape(host_shape);

  TF_RETURN_IF_ERROR(alias_config.ForEachAliasWithStatus(
      [&](const ShapeIndex& output_index,
          absl::optional<HloInputOutputAliasConfig::Alias> alias) {
        if (alias && alias->must_alias()) {
          VLOG(1) << alias->ToString();
          const MaybeOwningDeviceMemory& original_input =
              (*arguments)[alias->parameter_number].Buffers().element(
                  alias->parameter_index);
          if (!original_input.HasOwnership()) {
            return InvalidArgument(
                "An input was configured to be must-alias at "
                "compile time but not donated at runtime: %s",
                alias->ToString());
          }
        }
        return Status::OK();
      }));

  if (VLOG_IS_ON(3)) {
    VLOG(3) << "AllocateOutputMemoryWithInputReuse, device = " << device_ordinal
            << " host_shape = " << ShapeUtil::HumanStringWithLayout(host_shape);
    if (!Shape::Equal().MinorToMajorOnlyInLayout()(host_shape, device_shape)) {
      VLOG(3) << "Rewrote host_shape to device_shape: "
              << ShapeUtil::HumanStringWithLayout(host_shape) << " -> "
              << ShapeUtil::HumanStringWithLayout(device_shape);
    }
  }

  ExecutionOutput result(host_shape, std::move(device_shape), allocator,
                         device_ordinal);
  // Iterate through and allocate a buffer for each shape index, checking for
  // possible input buffer reuse.
  int64 reused_buffer_bytes = 0;
  int64 total_result_buffer_bytes = 0;
  for (auto& pair : result.MutableResult()->buffers()) {
    const ShapeIndex& result_index = pair.first;
    se::DeviceMemoryBase& result_buffer = pair.second;
    int64 allocation_bytes = ShapeSize(ShapeUtil::GetSubshape(
        result.Result().on_device_shape(), result_index));
    total_result_buffer_bytes += allocation_bytes;

    // Return an InternalError if result_index is invalid. This avoids failing
    // the CHECK when calling GetAliasedParameter
    if (!ShapeUtil::IndexIsValid(alias_config.shape(), result_index)) {
      return InternalError("result_index is invalid: %s",
                           result_index.ToString());
    }

    absl::optional<HloInputOutputAliasConfig::Alias> alias =
        alias_config.GetAliasedParameter(result_index);
    if (alias) {
      TF_RET_CHECK(alias->parameter_number < arguments->size());
      ExecutionInput& input = (*arguments)[alias->parameter_number];
      MaybeOwningDeviceMemory* device_memory =
          input.MutableBuffer(alias->parameter_index);
      if (auto owning = device_memory->Release()) {
        // If the caller passes the ownership of the device memory, reuse it
        // as the output buffer. It is up to the caller whether or not to
        // donate a buffer; the aliasing information describes which buffers
        // may alias, not buffers that must alias.
        se::DeviceMemoryBase device_memory_base = owning->Release();
        *device_memory = device_memory_base;
        result_buffer = device_memory_base;
        reused_buffer_bytes += allocation_bytes;
        // The caller is giving us the input buffer, but in case of error of the
        // execute call, we should not be releasing it as it contains valid data
        // (for example, it is a parameter which the user wants us to alias, in
        // a gradient update computation). So we store the index into the result
        // in the aliased vactor, which will be fed to the ExecutionOutput,
        // which will be using the indices to drop the addresses from its own
        // ScopedShapedBuffer result, if the ExecutionOutput is not committed.
        result.AddAliasedIndex(result_index);
      }
    }

    // We need to allocate a new output buffer for two cases:
    // - There is no alias between this output and any input.
    // - There is an alias, but the xla doesn't own the input memory so it can't
    // donate buffer to the computation.
    if (result_buffer.is_null()) {
      const Shape& on_device_shape = result.Result().on_device_shape();
      const Shape& on_device_subshape =
          ShapeUtil::GetSubshape(on_device_shape, result_index);
      TF_ASSIGN_OR_RETURN(
          auto allocated_buffer,
          allocator->Allocate(device_ordinal, allocation_bytes,
                              /*retry_on_failure=*/true,
                              on_device_subshape.layout().memory_space()));
      // Store the allocated buffer in our ScopedShapedBuffer, which takes
      // ownership.
      result_buffer = allocated_buffer.Release();
    }
    TF_RET_CHECK(allocation_bytes == 0 || result_buffer != nullptr);
  }

  VLOG(1) << "Reused " << reused_buffer_bytes
          << " parameter buffers (total result buffer size: "
          << total_result_buffer_bytes << ")";

  TF_RETURN_IF_ERROR(
      PopulateResultTupleBuffers(result.Result(), stream, transfer_stream));
  return std::move(result);
}

StatusOr<ExecutionOutput> TpuExecutableInterface::ExecuteAsyncOnStream(
    const ServiceExecutableRunOptions* run_options,
    std::vector<ExecutionInput> arguments,
    HloExecutionProfile* /*hlo_execution_profile*/) {
  std::vector<se::DeviceMemoryBase> memory_bases;
  memory_bases.reserve(arguments.size());
  for (auto& argument : arguments) {
    memory_bases.push_back(argument.Buffer({}).AsDeviceMemoryBase());
  }
  se::Stream* stream = run_options->stream();

  CHECK_NE(run_options->allocator(), nullptr);
  const Shape& shape =
      hlo_module_ == nullptr ? ShapeUtil::MakeNil() : result_shape();
  const HloInputOutputAliasConfig& alias_config =
      hlo_module_ == nullptr ? HloInputOutputAliasConfig()
                             : hlo_module_->input_output_alias_config();
  TF_ASSIGN_OR_RETURN(
      ExecutionOutput result,
      AllocateOutputMemoryWithInputReuse(
          shape, alias_config, run_options->allocator(), &arguments, stream,
          run_options->run_options().host_to_device_stream()));

  MarkToBeReleasedArguments(absl::MakeSpan(arguments), result);

  // Address of the buffer in TPU memory that is being speculated.
  absl::optional<se::DeviceMemoryBase> cross_program_prefetch_addr;
  if (hlo_module_) {
    for (const auto& [parameter, index] :
         hlo_module_->CrossProgramPrefetches()) {
      CHECK_LT(parameter, arguments.size());
      // Ensure the cross program prefetched buffer doesn't alias with any
      // program outputs. If the input and output aliased, the buffer could be
      // invalidated during program execution and the program could read stale
      // data from fast memory instead of fresh data in large memory.
      auto it = arguments[parameter].MutableBuffers()->find({index});
      CHECK(it != arguments[parameter].MutableBuffers()->end());
      if (absl::c_none_of(result.Result().buffers(), [&](auto index_addr_pair) {
            return index_addr_pair.second.IsSameAs(
                it->second.AsDeviceMemoryBase());
          })) {
        // Supports only one cross-program prefetch address.
        cross_program_prefetch_addr = it->second.AsDeviceMemoryBase();
      }
    }
  }

  TF_RETURN_IF_ERROR(LoadProgramAndEnqueueToStream(
      *run_options, memory_bases, result.Result().root_buffer(),
      cross_program_prefetch_addr));
  return std::move(result);
}

}  // namespace xla
<commit_msg>Fix C++ code not compatible with C++14 in tpu_executable_interface.cc<commit_after>/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/stream_executor/tpu/tpu_executable_interface.h"

#include <utility>

#include "absl/algorithm/container.h"
#include "tensorflow/compiler/xla/service/maybe_owning_device_memory.h"
#include "tensorflow/compiler/xla/service/shaped_buffer.h"
#include "tensorflow/compiler/xla/service/transfer_manager.h"
#include "tensorflow/compiler/xla/shape_util.h"
#include "tensorflow/compiler/xla/status_macros.h"
#include "tensorflow/compiler/xla/util.h"

namespace xla {

namespace {

// Write the tuple index buffers (arrays of pointers).
static Status PopulateResultTupleBuffers(const ShapedBuffer& result,
                                         se::Stream* stream,
                                         se::Stream* transfer_stream) {
  TF_ASSIGN_OR_RETURN(auto transfer_manager, TransferManager::GetForPlatform(
                                                 stream->parent()->platform()));
  if (transfer_manager->CanShapedBufferBeAccessedNow(stream->parent(),
                                                     result)) {
    TF_RETURN_IF_ERROR(transfer_manager->WriteTupleIndexTablesAsync(
        transfer_stream ? transfer_stream : stream, result));
    if (transfer_stream && transfer_stream != stream) {
      stream->ThenWaitFor(transfer_stream);
    }
    return Status::OK();
  } else {
    return transfer_manager->WriteTupleIndexTablesAsync(stream, result);
  }
}

}  // namespace

StatusOr<ExecutionOutput>
TpuExecutableInterface::AllocateOutputMemoryWithInputReuse(
    const Shape& host_shape, const HloInputOutputAliasConfig& alias_config,
    se::DeviceMemoryAllocator* allocator,
    std::vector<ExecutionInput>* arguments, se::Stream* stream,
    se::Stream* transfer_stream) {
  auto stream_exec = stream->parent();
  auto device_ordinal = stream_exec->device_ordinal();
  VLOG(3) << "AllocateOutputMemoryWithInputReuse, device = " << device_ordinal
          << " host_shape = " << ShapeUtil::HumanStringWithLayout(host_shape);
  Shape device_shape = HostShapeToDeviceShape(host_shape);

  TF_RETURN_IF_ERROR(alias_config.ForEachAliasWithStatus(
      [&](const ShapeIndex& output_index,
          absl::optional<HloInputOutputAliasConfig::Alias> alias) {
        if (alias && alias->must_alias()) {
          VLOG(1) << alias->ToString();
          const MaybeOwningDeviceMemory& original_input =
              (*arguments)[alias->parameter_number].Buffers().element(
                  alias->parameter_index);
          if (!original_input.HasOwnership()) {
            return InvalidArgument(
                "An input was configured to be must-alias at "
                "compile time but not donated at runtime: %s",
                alias->ToString());
          }
        }
        return Status::OK();
      }));

  if (VLOG_IS_ON(3)) {
    VLOG(3) << "AllocateOutputMemoryWithInputReuse, device = " << device_ordinal
            << " host_shape = " << ShapeUtil::HumanStringWithLayout(host_shape);
    if (!Shape::Equal().MinorToMajorOnlyInLayout()(host_shape, device_shape)) {
      VLOG(3) << "Rewrote host_shape to device_shape: "
              << ShapeUtil::HumanStringWithLayout(host_shape) << " -> "
              << ShapeUtil::HumanStringWithLayout(device_shape);
    }
  }

  ExecutionOutput result(host_shape, std::move(device_shape), allocator,
                         device_ordinal);
  // Iterate through and allocate a buffer for each shape index, checking for
  // possible input buffer reuse.
  int64 reused_buffer_bytes = 0;
  int64 total_result_buffer_bytes = 0;
  for (auto& pair : result.MutableResult()->buffers()) {
    const ShapeIndex& result_index = pair.first;
    se::DeviceMemoryBase& result_buffer = pair.second;
    int64 allocation_bytes = ShapeSize(ShapeUtil::GetSubshape(
        result.Result().on_device_shape(), result_index));
    total_result_buffer_bytes += allocation_bytes;

    // Return an InternalError if result_index is invalid. This avoids failing
    // the CHECK when calling GetAliasedParameter
    if (!ShapeUtil::IndexIsValid(alias_config.shape(), result_index)) {
      return InternalError("result_index is invalid: %s",
                           result_index.ToString());
    }

    absl::optional<HloInputOutputAliasConfig::Alias> alias =
        alias_config.GetAliasedParameter(result_index);
    if (alias) {
      TF_RET_CHECK(alias->parameter_number < arguments->size());
      ExecutionInput& input = (*arguments)[alias->parameter_number];
      MaybeOwningDeviceMemory* device_memory =
          input.MutableBuffer(alias->parameter_index);
      if (auto owning = device_memory->Release()) {
        // If the caller passes the ownership of the device memory, reuse it
        // as the output buffer. It is up to the caller whether or not to
        // donate a buffer; the aliasing information describes which buffers
        // may alias, not buffers that must alias.
        se::DeviceMemoryBase device_memory_base = owning->Release();
        *device_memory = device_memory_base;
        result_buffer = device_memory_base;
        reused_buffer_bytes += allocation_bytes;
        // The caller is giving us the input buffer, but in case of error of the
        // execute call, we should not be releasing it as it contains valid data
        // (for example, it is a parameter which the user wants us to alias, in
        // a gradient update computation). So we store the index into the result
        // in the aliased vactor, which will be fed to the ExecutionOutput,
        // which will be using the indices to drop the addresses from its own
        // ScopedShapedBuffer result, if the ExecutionOutput is not committed.
        result.AddAliasedIndex(result_index);
      }
    }

    // We need to allocate a new output buffer for two cases:
    // - There is no alias between this output and any input.
    // - There is an alias, but the xla doesn't own the input memory so it can't
    // donate buffer to the computation.
    if (result_buffer.is_null()) {
      const Shape& on_device_shape = result.Result().on_device_shape();
      const Shape& on_device_subshape =
          ShapeUtil::GetSubshape(on_device_shape, result_index);
      TF_ASSIGN_OR_RETURN(
          auto allocated_buffer,
          allocator->Allocate(device_ordinal, allocation_bytes,
                              /*retry_on_failure=*/true,
                              on_device_subshape.layout().memory_space()));
      // Store the allocated buffer in our ScopedShapedBuffer, which takes
      // ownership.
      result_buffer = allocated_buffer.Release();
    }
    TF_RET_CHECK(allocation_bytes == 0 || result_buffer != nullptr);
  }

  VLOG(1) << "Reused " << reused_buffer_bytes
          << " parameter buffers (total result buffer size: "
          << total_result_buffer_bytes << ")";

  TF_RETURN_IF_ERROR(
      PopulateResultTupleBuffers(result.Result(), stream, transfer_stream));
  return std::move(result);
}

StatusOr<ExecutionOutput> TpuExecutableInterface::ExecuteAsyncOnStream(
    const ServiceExecutableRunOptions* run_options,
    std::vector<ExecutionInput> arguments,
    HloExecutionProfile* /*hlo_execution_profile*/) {
  std::vector<se::DeviceMemoryBase> memory_bases;
  memory_bases.reserve(arguments.size());
  for (auto& argument : arguments) {
    memory_bases.push_back(argument.Buffer({}).AsDeviceMemoryBase());
  }
  se::Stream* stream = run_options->stream();

  CHECK_NE(run_options->allocator(), nullptr);
  const Shape& shape =
      hlo_module_ == nullptr ? ShapeUtil::MakeNil() : result_shape();
  const HloInputOutputAliasConfig& alias_config =
      hlo_module_ == nullptr ? HloInputOutputAliasConfig()
                             : hlo_module_->input_output_alias_config();
  TF_ASSIGN_OR_RETURN(
      ExecutionOutput result,
      AllocateOutputMemoryWithInputReuse(
          shape, alias_config, run_options->allocator(), &arguments, stream,
          run_options->run_options().host_to_device_stream()));

  MarkToBeReleasedArguments(absl::MakeSpan(arguments), result);

  // Address of the buffer in TPU memory that is being speculated.
  absl::optional<se::DeviceMemoryBase> cross_program_prefetch_addr;
  if (hlo_module_) {
    for (const auto& prefetch : hlo_module_->CrossProgramPrefetches()) {
      const auto& parameter = prefetch.first;
      const auto& index = prefetch.second;
      CHECK_LT(parameter, arguments.size());
      // Ensure the cross program prefetched buffer doesn't alias with any
      // program outputs. If the input and output aliased, the buffer could be
      // invalidated during program execution and the program could read stale
      // data from fast memory instead of fresh data in large memory.
      auto it = arguments[parameter].MutableBuffers()->find({index});
      CHECK(it != arguments[parameter].MutableBuffers()->end());
      if (absl::c_none_of(result.Result().buffers(), [&](auto index_addr_pair) {
            return index_addr_pair.second.IsSameAs(
                it->second.AsDeviceMemoryBase());
          })) {
        // Supports only one cross-program prefetch address.
        cross_program_prefetch_addr = it->second.AsDeviceMemoryBase();
      }
    }
  }

  TF_RETURN_IF_ERROR(LoadProgramAndEnqueueToStream(
      *run_options, memory_bases, result.Result().root_buffer(),
      cross_program_prefetch_addr));
  return std::move(result);
}

}  // namespace xla
<|endoftext|>


<commit_before>/**
 * @doc NIF helpers
 *
 * @copyright 2013 David H. Bronke
 * Licensed under the MIT license; see the LICENSE file for details.
 */

#include "nif_helpers.h"

bool getNIFDouble(ErlNifEnv* env, const ERL_NIF_TERM term, double* target)
{
	ulong ulongOther;
	long longOther;
	ErlNifUInt64 uint64Other;
	ErlNifSInt64 int64Other;

	if(!enif_get_double(env, term, target))
	{
		// Couldn't get a double directly; try the integer types...
		if(enif_get_ulong(env, term, &ulongOther))
		{
			*target = ulongOther;
		}
		else if(enif_get_long(env, term, &longOther))
		{
			*target = longOther;
		}
		else if(enif_get_uint64(env, term, &uint64Other))
		{
			*target = uint64Other;
		}
		else if(enif_get_int64(env, term, &int64Other))
		{
			*target = int64Other;
		}
		else
		{
			// Nothing matched what we were expecting; fail.
			return false;
		} // end if
	} // end if

	return true;
} // end getNIFDouble
<commit_msg>Corrected 'ulong' to 'unsigned long'.<commit_after>/**
 * @doc NIF helpers
 *
 * @copyright 2013 David H. Bronke
 * Licensed under the MIT license; see the LICENSE file for details.
 */

#include "nif_helpers.h"

bool getNIFDouble(ErlNifEnv* env, const ERL_NIF_TERM term, double* target)
{
	unsigned long ulongOther;
	long longOther;
	ErlNifUInt64 uint64Other;
	ErlNifSInt64 int64Other;

	if(!enif_get_double(env, term, target))
	{
		// Couldn't get a double directly; try the integer types...
		if(enif_get_ulong(env, term, &ulongOther))
		{
			*target = ulongOther;
		}
		else if(enif_get_long(env, term, &longOther))
		{
			*target = longOther;
		}
		else if(enif_get_uint64(env, term, &uint64Other))
		{
			*target = uint64Other;
		}
		else if(enif_get_int64(env, term, &int64Other))
		{
			*target = int64Other;
		}
		else
		{
			// Nothing matched what we were expecting; fail.
			return false;
		} // end if
	} // end if

	return true;
} // end getNIFDouble
<|endoftext|>