<commit_before>/*
 * The Biomechanical ToolKit
 * Copyright (c) 2009-2010, Arnaud Barré
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 *   * Redistributions of source code must retain the above
 *     copyright notice, this list of conditions and the following
 *     disclaimer.
 *   * Redistributions in binary form must reproduce the above
 *     copyright notice, this list of conditions and the following
 *     disclaimer in the documentation and/or other materials
 *     provided with the distribution.
 *   * Neither the name(s) of the copyright holders nor the names
 *     of its contributors may be used to endorse or promote products
 *     derived from this software without specific prior written
 *     permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 * COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
 * ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#include "btkBinaryFileStream_mmfstream_p.h"
#include "btkMacro.h"

#if defined(HAVE_SYS_MMAP)
  #if defined(HAVE_64_BIT)
    #define _LARGEFILE_SOURCE
    #define _FILE_OFFSET_BITS 64 // Replace all the functions (mmap, lssek, ...) by their 64 bit version
  #endif
  #include <sys/mman.h> // mmap, munmap
  #include <sys/stat.h> // fstat
  #include <fcntl.h> // open, close
#endif

namespace btk
{
  // ------------------------------------------------------------ //
  //                 Memory Mapped File Buffer                    //
  // ------------------------------------------------------------ //
  mmfilebuf* mmfilebuf::open(const char* s, std::ios_base::openmode mode)
  {
    if (this->is_open())
      return 0;
    std::ios_base::openmode mode_ = mode & (~std::ios_base::ate & ~std::ios_base::binary);
    if (mode_ & std::ios_base::out)
      this->m_Writing = true;
    else
      this->m_Writing = false;
    // Open the file and map it into the memory
    // The flags' extraction is inspired by the file fstream.cxx from the Comeau Computing library
#if defined(HAVE_SYS_MMAP) // POSIX
    // Select the flags for the function open
    int flags = 0;
    switch (mode_)
    {
    case std::ios_base::out:
    case std::ios_base::out + std::ios_base::trunc:
      flags = O_RDWR | O_CREAT | O_TRUNC;
      break;
    case std::ios_base::out + std::ios_base::app:
      flags = O_WRONLY | O_CREAT | O_APPEND;
      break;
    case std::ios_base::in:
      flags = O_RDONLY;
      break;
    case std::ios_base::in + std::ios_base::out:
      flags = O_RDWR;
      break;
    case std::ios_base::in + std::ios_base::out + std::ios_base::trunc:
      flags = O_RDWR | O_CREAT | O_TRUNC;
      break;
    default: // Other flags are not supported in the C++ standard
      return 0;
    }
    // Open the file
    if ((this->m_File = ::open(s, flags, S_IRWXU)) == -1)
      return 0;
    // Get the size of the file
    struct stat info;
    if (::fstat(this->m_File, &info) == -1)
      return this->close();
    this->m_BufferSize = this->m_LogicalSize = info.st_size;
    // New file or truncated file?
    if ((this->m_BufferSize == 0) && this->m_Writing)
    {
      this->m_BufferSize = mmfilebuf::granularity();
      if ((::lseek(this->m_File, this->m_BufferSize-1, SEEK_SET) == -1)
           || (::write(this->m_File, "", 1) == -1))
        return this->close();
    }
#else // Windows
    // Select the flags for the function CreateFile
    DWORD dwDesiredAccess, dwShareMode, dwCreationDisposition, dwFlagsandAttributes;
    dwDesiredAccess = GENERIC_READ | GENERIC_WRITE;
    dwShareMode = FILE_SHARE_READ;
    dwFlagsandAttributes = FILE_ATTRIBUTE_TEMPORARY;
    switch(mode_)
    {
    case std::ios_base::out:
      dwCreationDisposition = OPEN_ALWAYS;
      break;
    case std::ios_base::out + std::ios_base::trunc:
      dwCreationDisposition = CREATE_ALWAYS;
      break;
    case std::ios_base::out + std::ios_base::app:
      dwCreationDisposition = OPEN_ALWAYS;
      break;
    case std::ios_base::in:
      dwDesiredAccess = GENERIC_READ;
      dwShareMode = FILE_SHARE_READ | FILE_SHARE_WRITE; // Is it a good idea to give the possibility to other processes to write in the file that we are reading?
      dwCreationDisposition = OPEN_EXISTING;
      dwFlagsandAttributes = FILE_ATTRIBUTE_READONLY;
      break;
    case std::ios_base::in + std::ios_base::out:
      dwCreationDisposition = OPEN_ALWAYS;
      break;
    case std::ios_base::in + std::ios_base::out + std::ios_base::trunc:
      dwCreationDisposition = TRUNCATE_EXISTING;
      break;
    default: // Other flags are not supported in the C++ standard
      return 0;
    }
    // Open the file
    if ((this->m_File = ::CreateFile(s, dwDesiredAccess, dwShareMode, NULL,
                                     dwCreationDisposition, dwFlagsandAttributes, NULL))
        == INVALID_HANDLE_VALUE)
      return 0;
    // Get the size of the file
    DWORD dwFileSizeHi = 0;
    DWORD dwFileSizeLo = GetFileSize(this->m_File, &dwFileSizeHi);
    this->m_BufferSize = this->m_LogicalSize = (static_cast<std::streamsize>(dwFileSizeHi) << 32) | dwFileSizeLo;
    // New file or truncated file?
    if ((this->m_BufferSize == 0) && this->m_Writing)
    {
      this->m_BufferSize = mmfilebuf::granularity();
      LONG lDistHigh = this->m_BufferSize >> 32; // 32 = (sizeof(LONG) * 8)
      LONG lDistLow = this->m_BufferSize & 0xFFFFFFFF;
      if (((::SetFilePointer(this->m_File, lDistLow, &lDistHigh, FILE_BEGIN) == INVALID_SET_FILE_POINTER)
           && (::GetLastError() != NO_ERROR)) || (::SetEndOfFile(this->m_File) == 0))
        return this->close();
    }
#endif
    // Map the file
    if (!this->mapfile())
      return this->close();

    this->m_Position = 0;

    // If necessary go to the end of the file (::ate option)
    if ((mode & std::ios_base::ate) && (this->seekoff(0,std::ios_base::end, mode) == std::streampos(std::streamoff(-1))))
      return this->close();

    return this;
  };

  mmfilebuf* mmfilebuf::close()
  {
    if (!this->is_open())
      return 0;

#if defined(HAVE_SYS_MMAP)
    bool err = !(::munmap(this->mp_Buffer, this->m_BufferSize) == 0);
#else
    BOOL err = (::UnmapViewOfFile(this->mp_Buffer) == 0) || (::CloseHandle(this->m_Map) == 0);
    this->m_Map = NULL;
#endif

    // If in write mode, truncate the file to the number of bytes wrote.
    if (this->m_Writing)
    {
#if defined(HAVE_SYS_MMAP)
      err |= (::ftruncate(this->m_File, this->m_LogicalSize) == -1);
#else
      LONG lDistHigh = this->m_LogicalSize >> 32; // 32 = (sizeof(LONG) * 8)
      LONG lDistLow = this->m_LogicalSize & 0xFFFFFFFF;
      err |= (((::SetFilePointer(this->m_File, lDistLow, &lDistHigh, FILE_BEGIN) == INVALID_SET_FILE_POINTER)
               && (::GetLastError() != NO_ERROR)) || (::SetEndOfFile(this->m_File) == 0));
#endif
    }

    // Close the file
#if defined(HAVE_SYS_MMAP)
    err |= ::close(this->m_File);
#else
    err |= (::CloseHandle(this->m_File) == 0);
#endif
    this->m_File = BTK_MMFILEBUF_NO_FILE;

    if (err)
      return 0;
    this->mp_Buffer = 0;
    this->m_BufferSize = 0;
    this->m_LogicalSize = 0;

    return this;
  };

  std::streampos mmfilebuf::seekoff(std::streamoff off, std::ios_base::seekdir way, std::ios_base::openmode which)
  {
    btkNotUsed(which);
    switch(way)
    {
    case std::ios_base::beg:
      if (off < 0)
        return -1;
      this->m_Position = off;
      break;
    case std::ios_base::cur:
      if (this->m_Position + off < 0)
        return -1;
      this->m_Position += off;
      break;
    case std::ios_base::end:
      if (this->m_BufferSize + off < 0)
        return -1;
      this->m_Position = this->m_BufferSize + off;
      break;
    default:
      return -1;
    }
    return this->m_Position;
  };

  std::streampos mmfilebuf::seekpos(std::streampos pos, std::ios_base::openmode which)
  {
    btkNotUsed(which);
    if (pos < 0)
      return -1;

    this->m_Position = pos;
    return this->m_Position;
  };

  std::streamsize mmfilebuf::sgetn(char* s, std::streamsize n)
  {
    std::streamoff n_ = ((this->m_Position + n) > this->m_BufferSize) ? ((this->m_BufferSize - this->m_Position - 1) > 0 ? this->m_BufferSize - this->m_Position - 1 : 0) : n;
    for (std::streamoff i = 0 ; i < n_ ; ++i)
      s[i] = this->mp_Buffer[this->m_Position + i];
    this->m_Position += n_;
    return n_;
  };

  std::streamsize mmfilebuf::sputn(const char* s, std::streamsize n)
  {
    while ((this->m_Position + n) >= this->m_BufferSize)
    {
      if (!this->resizemap())
        return 0;
    }

    for (std::streamoff i = 0 ; i < n ; ++i)
      this->mp_Buffer[this->m_Position + i] = s[i];
    this->m_Position += n;

    if (this->m_Position >= this->m_LogicalSize)
      this->m_LogicalSize = this->m_Position;
    return n;
  };

  mmfilebuf* mmfilebuf::mapfile()
  {
#if defined(HAVE_SYS_MMAP)
    bool err = ((this->mp_Buffer = (char*)::mmap(0,
                                                 this->m_BufferSize,
                                                 this->m_Writing ? (PROT_READ | PROT_WRITE) : PROT_READ,
                                                 MAP_SHARED, this->m_File, 0)) == MAP_FAILED);
#else
    BOOL err = ((this->m_Map = ::CreateFileMapping(this->m_File,
                                                   NULL,
                                                   this->m_Writing ? PAGE_READWRITE : PAGE_READONLY,
                                                   0,
                                                   0,
                                                   NULL)) == NULL);
    err |= ((this->mp_Buffer = (char*)::MapViewOfFile(this->m_Map,
                                                       this->m_Writing ? FILE_MAP_WRITE : FILE_MAP_READ,
                                                       0,
                                                       0,
                                                       this->m_BufferSize)) == NULL);
#endif
    return err ? 0 : this;
  };

  mmfilebuf* mmfilebuf::resizemap()
  {
    if (!this->is_open() || !this->m_Writing)
      return 0;
    std::streamsize newBufferSize = this->m_BufferSize + this->granularity();
#if defined(_MSC_VER)
    if ((::UnmapViewOfFile(this->mp_Buffer) == 0) || (::CloseHandle(this->m_Map) == 0))
      return 0;
    this->m_Map = NULL;
    LONG lDistHigh = newBufferSize >> 32; // 32 = (sizeof(LONG) * 8)
    LONG lDistLow = newBufferSize & 0xFFFFFFFF;
    if (((::SetFilePointer(this->m_File, lDistLow, &lDistHigh, FILE_BEGIN) == INVALID_SET_FILE_POINTER)
               && (::GetLastError() != NO_ERROR)) || (::SetEndOfFile(this->m_File) == 0))
      return 0;
#else
    if (ftruncate(this->m_File, newBufferSize) == -1)
      return 0;
#endif
    this->m_BufferSize = newBufferSize;
    return this->mapfile();
  };

  int mmfilebuf::granularity()
  {
#if defined(_MSC_VER)
    SYSTEM_INFO info;
    ::GetSystemInfo(&info);
    return static_cast<int>(info.dwAllocationGranularity);
#else
    return static_cast<int>(::sysconf(_SC_PAGESIZE));
#endif
  };

  // ------------------------------------------------------------ //
  //                 Memory Mapped File Stream                    //
  // ------------------------------------------------------------ //

  mmfstream& mmfstream::read(char* s, std::streamsize n)
  {
    if (this->m_Filebuf.sgetn(s,n) != n)
      this->setstate(std::ios_base::eofbit | std::ios_base::failbit);
    else if (this->m_Filebuf.is_eob())
      this->setstate(std::ios_base::eofbit);
    return *this;
  };

  mmfstream& mmfstream::write(const char* s, std::streamsize n)
  {
    if (!this->m_Filebuf.writemode()) // Read-only?
      this->setstate(std::ios_base::failbit);
    else if (this->m_Filebuf.sputn(s,n) != n)
      this->setstate(std::ios_base::failbit);
    return *this;
  };
}<commit_msg>[FIX] The use of the memory mapped filestream under Windows 32-bits crashed BTK.<commit_after>/*
 * The Biomechanical ToolKit
 * Copyright (c) 2009-2010, Arnaud Barré
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 *   * Redistributions of source code must retain the above
 *     copyright notice, this list of conditions and the following
 *     disclaimer.
 *   * Redistributions in binary form must reproduce the above
 *     copyright notice, this list of conditions and the following
 *     disclaimer in the documentation and/or other materials
 *     provided with the distribution.
 *   * Neither the name(s) of the copyright holders nor the names
 *     of its contributors may be used to endorse or promote products
 *     derived from this software without specific prior written
 *     permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 * COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
 * ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#include "btkBinaryFileStream_mmfstream_p.h"
#include "btkMacro.h"

#if defined(HAVE_SYS_MMAP)
  #if defined(HAVE_64_BIT)
    #define _LARGEFILE_SOURCE
    #define _FILE_OFFSET_BITS 64 // Replace all the functions (mmap, lssek, ...) by their 64 bit version
  #endif
  #include <sys/mman.h> // mmap, munmap
  #include <sys/stat.h> // fstat
  #include <fcntl.h> // open, close
#endif

namespace btk
{
  // ------------------------------------------------------------ //
  //                 Memory Mapped File Buffer                    //
  // ------------------------------------------------------------ //
  mmfilebuf* mmfilebuf::open(const char* s, std::ios_base::openmode mode)
  {
    if (this->is_open())
      return 0;
    std::ios_base::openmode mode_ = mode & (~std::ios_base::ate & ~std::ios_base::binary);
    if (mode_ & std::ios_base::out)
      this->m_Writing = true;
    else
      this->m_Writing = false;
    // Open the file and map it into the memory
    // The flags' extraction is inspired by the file fstream.cxx from the Comeau Computing library
#if defined(HAVE_SYS_MMAP) // POSIX
    // Select the flags for the function open
    int flags = 0;
    switch (mode_)
    {
    case std::ios_base::out:
    case std::ios_base::out + std::ios_base::trunc:
      flags = O_RDWR | O_CREAT | O_TRUNC;
      break;
    case std::ios_base::out + std::ios_base::app:
      flags = O_WRONLY | O_CREAT | O_APPEND;
      break;
    case std::ios_base::in:
      flags = O_RDONLY;
      break;
    case std::ios_base::in + std::ios_base::out:
      flags = O_RDWR;
      break;
    case std::ios_base::in + std::ios_base::out + std::ios_base::trunc:
      flags = O_RDWR | O_CREAT | O_TRUNC;
      break;
    default: // Other flags are not supported in the C++ standard
      return 0;
    }
    // Open the file
    if ((this->m_File = ::open(s, flags, S_IRWXU)) == -1)
      return 0;
    // Get the size of the file
    struct stat info;
    if (::fstat(this->m_File, &info) == -1)
      return this->close();
    this->m_BufferSize = this->m_LogicalSize = info.st_size;
    // New file or truncated file?
    if ((this->m_BufferSize == 0) && this->m_Writing)
    {
      this->m_BufferSize = mmfilebuf::granularity();
      if ((::lseek(this->m_File, this->m_BufferSize-1, SEEK_SET) == -1)
           || (::write(this->m_File, "", 1) == -1))
        return this->close();
    }
#else // Windows
    // Select the flags for the function CreateFile
    DWORD dwDesiredAccess, dwShareMode, dwCreationDisposition, dwFlagsandAttributes;
    dwDesiredAccess = GENERIC_READ | GENERIC_WRITE;
    dwShareMode = FILE_SHARE_READ;
    dwFlagsandAttributes = FILE_ATTRIBUTE_TEMPORARY;
    switch(mode_)
    {
    case std::ios_base::out:
      dwCreationDisposition = OPEN_ALWAYS;
      break;
    case std::ios_base::out + std::ios_base::trunc:
      dwCreationDisposition = CREATE_ALWAYS;
      break;
    case std::ios_base::out + std::ios_base::app:
      dwCreationDisposition = OPEN_ALWAYS;
      break;
    case std::ios_base::in:
      dwDesiredAccess = GENERIC_READ;
      dwShareMode = FILE_SHARE_READ | FILE_SHARE_WRITE; // Is it a good idea to give the possibility to other processes to write in the file that we are reading?
      dwCreationDisposition = OPEN_EXISTING;
      dwFlagsandAttributes = FILE_ATTRIBUTE_READONLY;
      break;
    case std::ios_base::in + std::ios_base::out:
      dwCreationDisposition = OPEN_ALWAYS;
      break;
    case std::ios_base::in + std::ios_base::out + std::ios_base::trunc:
      dwCreationDisposition = TRUNCATE_EXISTING;
      break;
    default: // Other flags are not supported in the C++ standard
      return 0;
    }
    // Open the file
    if ((this->m_File = ::CreateFile(s, dwDesiredAccess, dwShareMode, NULL,
                                     dwCreationDisposition, dwFlagsandAttributes, NULL))
        == INVALID_HANDLE_VALUE)
      return 0;
    // Get the size of the file
    DWORD dwFileSizeHi = 0;
    DWORD dwFileSizeLo = GetFileSize(this->m_File, &dwFileSizeHi);
    this->m_BufferSize = this->m_LogicalSize = static_cast<std::streamsize>(((uint64_t)dwFileSizeHi << 32) | dwFileSizeLo);
    // New file or truncated file?
    if ((this->m_BufferSize == 0) && this->m_Writing)
    {
      this->m_BufferSize = mmfilebuf::granularity();
      LONG lDistHigh = (uint64_t)this->m_BufferSize >> 32; // 32 = (sizeof(LONG) * 8)
      LONG lDistLow = this->m_BufferSize & 0xFFFFFFFF;
      if (((::SetFilePointer(this->m_File, lDistLow, &lDistHigh, FILE_BEGIN) == INVALID_SET_FILE_POINTER)
           && (::GetLastError() != NO_ERROR)) || (::SetEndOfFile(this->m_File) == 0))
        return this->close();
    }
#endif
    // Map the file
    if (!this->mapfile())
      return this->close();

    this->m_Position = 0;

    // If necessary go to the end of the file (::ate option)
    if ((mode & std::ios_base::ate) && (this->seekoff(0,std::ios_base::end, mode) == std::streampos(std::streamoff(-1))))
      return this->close();

    return this;
  };

  mmfilebuf* mmfilebuf::close()
  {
    if (!this->is_open())
      return 0;

#if defined(HAVE_SYS_MMAP)
    bool err = !(::munmap(this->mp_Buffer, this->m_BufferSize) == 0);
#else
    BOOL err = (::UnmapViewOfFile(this->mp_Buffer) == 0) || (::CloseHandle(this->m_Map) == 0);
    this->m_Map = NULL;
#endif

    // If in write mode, truncate the file to the number of bytes wrote.
    if (this->m_Writing)
    {
#if defined(HAVE_SYS_MMAP)
      err |= (::ftruncate(this->m_File, this->m_LogicalSize) == -1);
#else
      LONG lDistHigh = (uint64_t)this->m_LogicalSize >> 32; // 32 = (sizeof(LONG) * 8)
      LONG lDistLow = this->m_LogicalSize & 0xFFFFFFFF;
      err |= (((::SetFilePointer(this->m_File, lDistLow, &lDistHigh, FILE_BEGIN) == INVALID_SET_FILE_POINTER)
               && (::GetLastError() != NO_ERROR)) || (::SetEndOfFile(this->m_File) == 0));
#endif
    }

    // Close the file
#if defined(HAVE_SYS_MMAP)
    err |= ::close(this->m_File);
#else
    err |= (::CloseHandle(this->m_File) == 0);
#endif
    this->m_File = BTK_MMFILEBUF_NO_FILE;

    if (err)
      return 0;
    this->mp_Buffer = 0;
    this->m_BufferSize = 0;
    this->m_LogicalSize = 0;

    return this;
  };

  std::streampos mmfilebuf::seekoff(std::streamoff off, std::ios_base::seekdir way, std::ios_base::openmode which)
  {
    btkNotUsed(which);
    switch(way)
    {
    case std::ios_base::beg:
      if (off < 0)
        return -1;
      this->m_Position = off;
      break;
    case std::ios_base::cur:
      if (this->m_Position + off < 0)
        return -1;
      this->m_Position += off;
      break;
    case std::ios_base::end:
      if (this->m_BufferSize + off < 0)
        return -1;
      this->m_Position = this->m_BufferSize + off;
      break;
    default:
      return -1;
    }
    return this->m_Position;
  };

  std::streampos mmfilebuf::seekpos(std::streampos pos, std::ios_base::openmode which)
  {
    btkNotUsed(which);
    if (pos < 0)
      return -1;

    this->m_Position = pos;
    return this->m_Position;
  };

  std::streamsize mmfilebuf::sgetn(char* s, std::streamsize n)
  {
    std::streamoff n_ = ((this->m_Position + n) > this->m_BufferSize) ? ((this->m_BufferSize - this->m_Position - 1) > 0 ? this->m_BufferSize - this->m_Position - 1 : 0) : n;
    for (std::streamoff i = 0 ; i < n_ ; ++i)
      s[i] = this->mp_Buffer[this->m_Position + i];
    this->m_Position += n_;
    return n_;
  };

  std::streamsize mmfilebuf::sputn(const char* s, std::streamsize n)
  {
    while ((this->m_Position + n) >= this->m_BufferSize)
    {
      if (!this->resizemap())
        return 0;
    }

    for (std::streamoff i = 0 ; i < n ; ++i)
      this->mp_Buffer[this->m_Position + i] = s[i];
    this->m_Position += n;

    if (this->m_Position >= this->m_LogicalSize)
      this->m_LogicalSize = this->m_Position;
    return n;
  };

  mmfilebuf* mmfilebuf::mapfile()
  {
#if defined(HAVE_SYS_MMAP)
    bool err = ((this->mp_Buffer = (char*)::mmap(0,
                                                 this->m_BufferSize,
                                                 this->m_Writing ? (PROT_READ | PROT_WRITE) : PROT_READ,
                                                 MAP_SHARED, this->m_File, 0)) == MAP_FAILED);
#else
    BOOL err = ((this->m_Map = ::CreateFileMapping(this->m_File,
                                                   NULL,
                                                   this->m_Writing ? PAGE_READWRITE : PAGE_READONLY,
                                                   0,
                                                   0,
                                                   NULL)) == NULL);
    err |= ((this->mp_Buffer = (char*)::MapViewOfFile(this->m_Map,
                                                       this->m_Writing ? FILE_MAP_WRITE : FILE_MAP_READ,
                                                       0,
                                                       0,
                                                       this->m_BufferSize)) == NULL);
#endif
    return err ? 0 : this;
  };

  mmfilebuf* mmfilebuf::resizemap()
  {
    if (!this->is_open() || !this->m_Writing)
      return 0;
    std::streamsize newBufferSize = this->m_BufferSize + this->granularity();
#if defined(_MSC_VER)
    if ((::UnmapViewOfFile(this->mp_Buffer) == 0) || (::CloseHandle(this->m_Map) == 0))
      return 0;
    this->m_Map = NULL;
    LONG lDistHigh = (uint64_t)newBufferSize >> 32; // 32 = (sizeof(LONG) * 8)
    LONG lDistLow = newBufferSize & 0xFFFFFFFF;
    if (((::SetFilePointer(this->m_File, lDistLow, &lDistHigh, FILE_BEGIN) == INVALID_SET_FILE_POINTER)
               && (::GetLastError() != NO_ERROR)) || (::SetEndOfFile(this->m_File) == 0))
      return 0;
#else
    if (ftruncate(this->m_File, newBufferSize) == -1)
      return 0;
#endif
    this->m_BufferSize = newBufferSize;
    return this->mapfile();
  };

  int mmfilebuf::granularity()
  {
#if defined(_MSC_VER)
    SYSTEM_INFO info;
    ::GetSystemInfo(&info);
    return static_cast<int>(info.dwAllocationGranularity);
#else
    return static_cast<int>(::sysconf(_SC_PAGESIZE));
#endif
  };

  // ------------------------------------------------------------ //
  //                 Memory Mapped File Stream                    //
  // ------------------------------------------------------------ //

  mmfstream& mmfstream::read(char* s, std::streamsize n)
  {
    if (this->m_Filebuf.sgetn(s,n) != n)
      this->setstate(std::ios_base::eofbit | std::ios_base::failbit);
    else if (this->m_Filebuf.is_eob())
      this->setstate(std::ios_base::eofbit);
    return *this;
  };

  mmfstream& mmfstream::write(const char* s, std::streamsize n)
  {
    if (!this->m_Filebuf.writemode()) // Read-only?
      this->setstate(std::ios_base::failbit);
    else if (this->m_Filebuf.sputn(s,n) != n)
      this->setstate(std::ios_base::failbit);
    return *this;
  };
}<|endoftext|>


<commit_before>#include "otbVectorImage.h"
#include "otbImage.h"
#include "otbImageFileReader.h"
#include "otbStreamingImageFileWriter.h"
#include "otbImageFileWriter.h"
#include "otbSOMMap.h"
#include "otbSOM.h"
#include "otbSOMImageClassificationFilter.h"
#include "otbCommandLineArgumentParser.h"
#include "itkEuclideanDistance.h"
#include "itkImageRegionSplitter.h"
#include "otbStreamingTraits.h"
#include "itkImageRegionConstIterator.h"
#include "itkVariableSizeMatrix.h"
#include "itkListSample.h"
#include "itkVariableLengthVector.h"

int main(int argc, char * argv[])
{

  // Parse command line parameters
  typedef otb::CommandLineArgumentParser ParserType;
  ParserType::Pointer parser = ParserType::New();

  parser->SetProgramDescription("Unsupervised Self Organizing Map image classification");
  parser->AddInputImage();
  parser->AddOutputImage();
  parser->AddOption("--ValidityMask","Validity mask","-vm", 1, true);
  parser->AddOption("--MaxTrainingSetSize","Size of the training set","-ts", 1, true);
  parser->AddOption("--TrainingSetProbability","Probability for a sample to be selected in the training set","-tp", 1, true);
  parser->AddOption("--StreamingNumberOfLines","Number of lined for each streaming block","-sl", 1, true);
  parser->AddOption("--SOMMap","Output SOM map","-sm", 1, true);
  parser->AddOption("--SizeX","X size of the SOM map","-sx", 1, true);
  parser->AddOption("--SizeY","Y size of the SOM map","-sy", 1, true);
  parser->AddOption("--NeighborhoodInitX","X initial neighborhood of the SOM map","-nx", 1, true);
  parser->AddOption("--NeighborhoodInitY","Y initial neighborhood of the SOM map","-ny", 1, true);
  parser->AddOption("--NumberOfIterations","Number of iterations of the SOM learning","-ni", 1, true);
  parser->AddOption("--BetaInit","Initial beta value","-bi", 1, true);
  parser->AddOption("--BetaFinal","Final beta value","-be", 1, true);
  parser->AddOption("--InitValue","Initial value","-iv", 1, true);

  typedef otb::CommandLineArgumentParseResult ParserResultType;
  ParserResultType::Pointer  parseResult = ParserResultType::New();

  try
  {
    parser->ParseCommandLine(argc, argv, parseResult);
  }
  catch ( itk::ExceptionObject & err )
  {
    std::string descriptionException = err.GetDescription();
    if (descriptionException.find("ParseCommandLine(): Help Parser") != std::string::npos)
    {
      return EXIT_SUCCESS;
    }
    if (descriptionException.find("ParseCommandLine(): Version Parser") != std::string::npos)
    {
      return EXIT_SUCCESS;
    }
    return EXIT_FAILURE;
  }

  // initiating random number generation
  srand(time(NULL));

  std::string infname = parseResult->GetInputImage();
  std::string maskfname = parseResult->GetParameterString("--ValidityMask", 0);
  std::string outfname = parseResult->GetOutputImage();
  std::string somfname = parseResult->GetParameterString("--SOMMap", 0);
  const unsigned int nbsamples = parseResult->GetParameterUInt("--MaxTrainingSetSize");
  const double trainingProb =parseResult->GetParameterDouble("--TrainingSetProbability");
  const unsigned int nbLinesForStreaming = parseResult->GetParameterUInt("--StreamingNumberOfLines");
  const unsigned int sizeX = parseResult->GetParameterUInt("--SizeX");
  const unsigned int sizeY = parseResult->GetParameterUInt("--SizeY");
  const unsigned int neighInitX = parseResult->GetParameterUInt("--NeighborhoodInitX");
  const unsigned int neighInitY= parseResult->GetParameterUInt("--NeighborhoodInitY");
  const unsigned int nbIterations=parseResult->GetParameterUInt("--NumberOfIterations");
  const double betaInit = parseResult->GetParameterDouble("--BetaInit");
  const double betaEnd= parseResult->GetParameterDouble("--BetaFinal");
  const float initValue = parseResult->GetParameterFloat("--InitValue");

  typedef float PixelType;
  typedef unsigned short LabeledPixelType;

  typedef otb::VectorImage<PixelType, 2> ImageType;
  typedef otb::Image<LabeledPixelType, 2> LabeledImageType;
  typedef otb::ImageFileReader<ImageType> ImageReaderType;
  typedef otb::ImageFileReader<LabeledImageType> LabeledImageReaderType;
  typedef otb::StreamingImageFileWriter<LabeledImageType> WriterType;

  typedef itk::VariableLengthVector<double> SampleType;
  typedef itk::Statistics::EuclideanDistance<SampleType> DistanceType;
  typedef otb::SOMMap<SampleType, DistanceType, 2> SOMMapType;
  typedef itk::Statistics::ListSample<SampleType> ListSampleType;
  typedef otb::SOM<ListSampleType, SOMMapType> EstimatorType;
  typedef otb::ImageFileWriter<ImageType> SOMMapWriterType;

  typedef otb::StreamingTraits<ImageType> StreamingTraitsType;
  typedef itk::ImageRegionSplitter<2>  SplitterType;
  typedef ImageType::RegionType RegionType;

  typedef itk::ImageRegionConstIterator<ImageType> IteratorType;
  typedef itk::ImageRegionConstIterator<LabeledImageType> LabeledIteratorType;
  typedef itk::ImageRegionConstIterator<SOMMapType> SOMIteratorType;

  typedef otb::SOMImageClassificationFilter<ImageType, LabeledImageType, SOMMapType> ClassificationFilterType;

  ImageReaderType::Pointer reader = ImageReaderType::New();
  LabeledImageReaderType::Pointer maskReader = LabeledImageReaderType::New();

  reader->SetFileName(infname);
  maskReader->SetFileName(maskfname);

  /*******************************************/
  /*           Sampling data                 */
  /*******************************************/
  std::cout<<std::endl;
  std::cout<<"-- SAMPLING DATA --"<<std::endl;
  std::cout<<std::endl;

  // Update input images information
  reader->GenerateOutputInformation();
  maskReader->GenerateOutputInformation();

  if (reader->GetOutput()->GetLargestPossibleRegion()
      != maskReader->GetOutput()->GetLargestPossibleRegion()
     )
  {
    std::cerr<<"Mask image and input image have different sizes."<<std::endl;
    return EXIT_FAILURE;
  }

  RegionType largestRegion = reader->GetOutput()->GetLargestPossibleRegion();

  // Setting up local streaming capabilities
  SplitterType::Pointer splitter = SplitterType::New();
  unsigned int numberOfStreamDivisions = StreamingTraitsType::CalculateNumberOfStreamDivisions(reader->GetOutput(),
                                         largestRegion,
                                         splitter,
                                         otb::SET_BUFFER_NUMBER_OF_LINES,
                                         0, 0, nbLinesForStreaming);

  std::cout<<"The images will be streamed into "<<numberOfStreamDivisions<<" parts."<<std::endl;

  // Training sample lists
  ListSampleType::Pointer sampleList = ListSampleType::New();

  // Sample dimension and max dimension
  unsigned int sampleSize = reader->GetOutput()->GetNumberOfComponentsPerPixel();
  unsigned int totalSamples = 0;
  std::cout<<"The following sample size will be used: "<<sampleSize<<std::endl;
  std::cout<<std::endl;
  // local streaming variables
  unsigned int piece = 0;
  RegionType streamingRegion;

  while (totalSamples<nbsamples)
  {
    piece = static_cast<unsigned int>(static_cast<double>(numberOfStreamDivisions)*rand()/(RAND_MAX));

    streamingRegion = splitter->GetSplit(piece, numberOfStreamDivisions, largestRegion);

    std::cout<<"Processing region: "<<streamingRegion<<std::endl;

    reader->GetOutput()->SetRequestedRegion(streamingRegion);
    reader->GetOutput()->PropagateRequestedRegion();
    reader->GetOutput()->UpdateOutputData();

    maskReader->GetOutput()->SetRequestedRegion(streamingRegion);
    maskReader->GetOutput()->PropagateRequestedRegion();
    maskReader->GetOutput()->UpdateOutputData();

    IteratorType it(reader->GetOutput(), streamingRegion);
    LabeledIteratorType maskIt(maskReader->GetOutput(), streamingRegion);

    it.GoToBegin();
    maskIt.GoToBegin();

    unsigned int localNbSamples=0;

    // Loop on the image
    while (!it.IsAtEnd()&&!maskIt.IsAtEnd()&&(totalSamples<nbsamples))
    {
      // If the current pixel is labeled
      if (maskIt.Get()>0)
      {
        if ((rand()<trainingProb*RAND_MAX))
        {
          SampleType newSample;
          newSample.SetSize(sampleSize);
          // build the sample
          newSample.Fill(0);
          for (unsigned int i = 0; i<sampleSize; ++i)
          {
            newSample[i]=it.Get()[i];
          }
          // Update the the sample lists
          sampleList->PushBack(newSample);
          ++totalSamples;
          ++localNbSamples;
        }
      }
      ++it;
      ++maskIt;
    }
    std::cout<<localNbSamples<<" samples added to the training set."<<std::endl;
    std::cout<<std::endl;
  }

  std::cout<<"The final training set contains "<<totalSamples<<" samples."<<std::endl;

  std::cout<<std::endl;
  std::cout<<"Data sampling completed."<<std::endl;
  std::cout<<std::endl;

  /*******************************************/
  /*           Learning                      */
  /*******************************************/

  std::cout<<"-- LEARNING --"<<std::endl;

  EstimatorType::Pointer estimator = EstimatorType::New();

  estimator->SetListSample(sampleList);
  EstimatorType::SizeType size;
  size[0]=sizeX;
  size[1]=sizeY;
  estimator->SetMapSize(size);
  EstimatorType::SizeType radius;
  radius[0] = neighInitX;
  radius[1] = neighInitY;
  estimator->SetNeighborhoodSizeInit(radius);
  estimator->SetNumberOfIterations(nbIterations);
  estimator->SetBetaInit(betaInit);
  estimator->SetBetaEnd(betaEnd);
  estimator->SetMaxWeight(initValue);
//  estimator->SetRandomInit(true);
//   estimator->SetSeed(time(NULL));
  estimator->Update();

  ImageType::Pointer vectormap = ImageType::New();
  vectormap->SetRegions(estimator->GetOutput()->GetLargestPossibleRegion());
  vectormap->SetNumberOfComponentsPerPixel(108);
  vectormap->Allocate();

  ImageType::PixelType black;
  black.SetSize(108);
  black.Fill(0);
  vectormap->FillBuffer(black);

  SOMIteratorType somIt(estimator->GetOutput(), estimator->GetOutput()->GetLargestPossibleRegion());
  IteratorType vectorIt(vectormap, estimator->GetOutput()->GetLargestPossibleRegion());

  somIt.GoToBegin();
  vectorIt.GoToBegin();

  while (!somIt.IsAtEnd() && !vectorIt.IsAtEnd())
  {
    for (unsigned int i = 0; i<somIt.Get().GetSize(); ++i)
    {
      vectorIt.Get()[i]=somIt.Get()[i];
    }
    ++somIt;
    ++vectorIt;
  }

  SOMMapWriterType::Pointer somWriter = SOMMapWriterType::New();
  somWriter->SetFileName(somfname);
  somWriter->SetInput(vectormap);
  somWriter->Update();

  std::cout<<std::endl;
  std::cout<<"Learning completed."<<std::endl;
  std::cout<<std::endl;

  /*******************************************/
  /*           Classification                */
  /*******************************************/

  std::cout<<"-- CLASSIFICATION --"<<std::endl;
  std::cout<<std::endl;
  ClassificationFilterType::Pointer classifier = ClassificationFilterType::New();
  classifier->SetInput(reader->GetOutput());
  classifier->SetInputMask(maskReader->GetOutput());
  classifier->SetMap(estimator->GetOutput());

  WriterType::Pointer writer = WriterType::New();
  writer->SetFileName(outfname);
  writer->SetInput(classifier->GetOutput());
  writer->SetNumberOfDivisionsStrippedStreaming(numberOfStreamDivisions);
  writer->Update();

  std::cout<<"Classification completed."<<std::endl;
  std::cout<<std::endl;

  return EXIT_SUCCESS;
}
<commit_msg>ENH: adapt SOMClassification to new framework<commit_after>/*=========================================================================

 Program:   ORFEO Toolbox
 Language:  C++
 Date:      $Date$
 Version:   $Revision$

 Copyright (c) Centre National d'Etudes Spatiales. All rights reserved.
 See OTBCopyright.txt for details.

 This software is distributed WITHOUT ANY WARRANTY; without even
 the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
 PURPOSE.  See the above copyright notices for more information.

 =========================================================================*/
#include "otbWrapperApplication.h"
#include "otbWrapperApplicationFactory.h"

#include "otbSOMMap.h"
#include "otbSOM.h"
#include "otbSOMImageClassificationFilter.h"
#include "itkEuclideanDistance.h"
#include "itkImageRegionSplitter.h"
#include "otbStreamingTraits.h"
#include "itkImageRegionConstIterator.h"
// #include "itkListSample.h"
#include "itkImageRandomNonRepeatingConstIteratorWithIndex.h"
//#include "itkMersenneTwisterRandomVariateGenerator.h"

namespace otb
{
namespace Wrapper
{

class SOMClassification : public Application
{
public:
  /** Standard class typedefs. */
  typedef SOMClassification             Self;
  typedef Application                   Superclass;
  typedef itk::SmartPointer<Self>       Pointer;
  typedef itk::SmartPointer<const Self> ConstPointer;

  /** Standard macro */
  itkNewMacro(Self);

  itkTypeMacro(SOMClassification, otb::Application);

  /** Filters typedef */
  typedef UInt16ImageType LabeledImageType;

  typedef itk::VariableLengthVector<double> SampleType;
  typedef itk::Statistics::EuclideanDistance<SampleType> DistanceType;
  typedef otb::SOMMap<SampleType, DistanceType, 2> SOMMapType;
  typedef itk::Statistics::ListSample<SampleType> ListSampleType;
  typedef otb::SOM<ListSampleType, SOMMapType> EstimatorType;

  typedef otb::StreamingTraits<FloatVectorImageType> StreamingTraitsType;
  typedef itk::ImageRegionSplitter<2>  SplitterType;
  typedef FloatVectorImageType::RegionType RegionType;

  typedef itk::ImageRegionConstIterator<FloatVectorImageType> IteratorType;
  typedef itk::ImageRegionConstIterator<LabeledImageType> LabeledIteratorType;
  typedef itk::ImageRegionConstIterator<SOMMapType> SOMIteratorType;

  typedef otb::SOMImageClassificationFilter
          <FloatVectorImageType, LabeledImageType, SOMMapType> ClassificationFilterType;

private:
  SOMClassification()
  {
    SetName("SOMClassification");
    SetDescription("SOM image classification.");

    // Documentation
    SetDocName("SOM Classification Application");
    SetDocLongDescription("Unsupervised Self Organizing Map image classification.");
    SetDocLimitations("None");
    SetDocAuthors("OTB-Team");
    SetDocSeeAlso(" ");

    AddDocTag(Tags::Segmentation);
    AddDocTag(Tags::Learning);
  }

  virtual ~SOMClassification()
  {
  }

  void DoCreateParameters()
    {
    AddParameter(ParameterType_InputImage,  "in",   "InputImage");
    SetParameterDescription("in", "Input image.");

    AddParameter(ParameterType_OutputImage,  "out",   "OutputImage");
    SetParameterDescription("out", "Output classified image.");

    AddParameter(ParameterType_InputImage,  "vm",   "ValidityMask");
    SetParameterDescription("vm", "Validity mask");

    AddParameter(ParameterType_Float, "tp", "TrainingProbability");
    SetParameterDescription("tp", "Probability for a sample to be selected in the training set");

    AddParameter(ParameterType_Int,  "ts",   "TrainingSet");
    SetParameterDescription("ts", "Maximum training set size");

    AddParameter(ParameterType_Int,  "sl",   "StreamingLines");
    SetParameterDescription("sl", "Number of lines in each streaming block (used during data sampling)");

    AddParameter(ParameterType_OutputImage, "som", "SOM Map");
    SetParameterDescription("som","Self-Organizing Map map");

    AddParameter(ParameterType_Int,  "sx",   "SizeX");
    SetParameterDescription("sx", "X size of the SOM map");

    AddParameter(ParameterType_Int,  "sy",   "SizeY");
    SetParameterDescription("sy", "Y size of the SOM map");

    AddParameter(ParameterType_Int,  "nx",   "NeighborhoodX");
    SetParameterDescription("nx", "X initial neighborhood of the SOM map");

    AddParameter(ParameterType_Int,  "ny",   "NeighborhoodY");
    SetParameterDescription("ny", "Y initial neighborhood of the SOM map");

    AddParameter(ParameterType_Int,  "ni",   "NumberIteration");
    SetParameterDescription("ni", "Number of iterations of the SOM learning");

    AddParameter(ParameterType_Float,  "bi",   "BetaInit");
    SetParameterDescription("bi", "Initial beta value");

    AddParameter(ParameterType_Float,  "bf",   "BetaFinal");
    SetParameterDescription("bf", "Final beta value");

    AddParameter(ParameterType_Float,  "iv",   "InitialValue");
    SetParameterDescription("iv", "Initial value (max weight)");

    AddParameter(ParameterType_RAM, "ram", "Available RAM");
    SetDefaultParameterInt("ram", 256);
    MandatoryOff("ram");
    // TODO : replace StreamingLines by RAM param ?

    // Default parameters
    SetDefaultParameterInt("sx", 32);
    SetDefaultParameterInt("sy", 32);
    SetDefaultParameterInt("nx", 10);
    SetDefaultParameterInt("ny", 10);
    SetDefaultParameterInt("ni", 5);
    SetDefaultParameterFloat("bi",1.0);
    SetDefaultParameterFloat("bf",0.1);
    SetDefaultParameterFloat("iv",0.0);

    // Doc example parameter settings
    SetDocExampleParameterValue("in", "poupees_sub.png");
    SetDocExampleParameterValue("out","poupees_classif.tif");
    SetDocExampleParameterValue("vm", "BASELINE/leSOMPoupeesClassified.hdr");
    SetDocExampleParameterValue("tp", "1.0");
    SetDocExampleParameterValue("ts","16384");
    SetDocExampleParameterValue("sl", "32");
    SetDocExampleParameterValue("som", "poupees_map.hdr");
    SetDocExampleParameterValue("sx", "32");
    SetDocExampleParameterValue("sy", "32");
    SetDocExampleParameterValue("nx", "10");
    SetDocExampleParameterValue("ny", "10");
    SetDocExampleParameterValue("ni", "5");
    SetDocExampleParameterValue("bi", "1.0");
    SetDocExampleParameterValue("bf", "0.1");
    SetDocExampleParameterValue("iv", "0");
    }

  void DoUpdateParameters()
    {
    // Nothing to do
    }

  void DoExecute()
    {
    // initiating random number generation
    itk::Statistics::MersenneTwisterRandomVariateGenerator::Pointer
        randomGen = itk::Statistics::MersenneTwisterRandomVariateGenerator::New();

    FloatVectorImageType::Pointer input = GetParameterImage("in");
    LabeledImageType::Pointer mask = GetParameterImage<LabeledImageType>("vm");

    /*******************************************/
    /*           Sampling data                 */
    /*******************************************/
    otbAppLogINFO("-- SAMPLING DATA --");

    if (input->GetLargestPossibleRegion()
        != mask->GetLargestPossibleRegion())
    {
      otbAppLogFATAL("Mask image and input image have different sizes.");
    }

    RegionType largestRegion = input->GetLargestPossibleRegion();

    // Setting up local streaming capabilities
    SplitterType::Pointer splitter = SplitterType::New();
    unsigned int numberOfStreamDivisions = StreamingTraitsType::CalculateNumberOfStreamDivisions(input,
                                          largestRegion,
                                          splitter,
                                          otb::SET_BUFFER_NUMBER_OF_LINES,
                                          0, 0, GetParameterInt("sl"));

    otbAppLogINFO("The images will be streamed into "<<numberOfStreamDivisions<<" parts.");

    // Training sample lists
    ListSampleType::Pointer sampleList = ListSampleType::New();
    const double trainingProb = static_cast<double>(GetParameterFloat("tp"));
    const unsigned int nbsamples = GetParameterInt("ts");

    // Sample dimension and max dimension
    unsigned int sampleSize = input->GetNumberOfComponentsPerPixel();
    unsigned int totalSamples = 0;
    otbAppLogINFO("The following sample size will be used: "<<sampleSize);

    // local streaming variables
    unsigned int piece = 0;
    RegionType streamingRegion;

    // create a random permutation to explore
    itk::RandomPermutation randPerm(numberOfStreamDivisions);
    unsigned int index = 0;

    // TODO : maybe change the approach: at the moment, the sampling process is able to pick a sample twice or more
    while (totalSamples < nbsamples)
    {
      piece = randPerm[index];

      streamingRegion = splitter->GetSplit(piece, numberOfStreamDivisions, largestRegion);

      //otbAppLogINFO("Processing region: "<<streamingRegion);

      input->SetRequestedRegion(streamingRegion);
      input->PropagateRequestedRegion();
      input->UpdateOutputData();

      mask->SetRequestedRegion(streamingRegion);
      mask->PropagateRequestedRegion();
      mask->UpdateOutputData();

      IteratorType it(input, streamingRegion);
      LabeledIteratorType maskIt(mask, streamingRegion);

      it.GoToBegin();
      maskIt.GoToBegin();

      unsigned int localNbSamples=0;

      // Loop on the image
      while ( !it.IsAtEnd()
           && !maskIt.IsAtEnd()
           && (totalSamples<nbsamples))
      {
        // If the current pixel is labeled
        if (maskIt.Get()>0)
        {
          if (randomGen->GetVariateWithClosedRange() < trainingProb)
          {
            SampleType newSample;
            newSample.SetSize(sampleSize);
            // build the sample
            newSample.Fill(0);
            for (unsigned int i = 0; i<sampleSize; ++i)
            {
              newSample[i]=it.Get()[i];
            }
            // Update the sample lists
            sampleList->PushBack(newSample);
            ++totalSamples;
            ++localNbSamples;
          }
        }
        ++it;
        ++maskIt;
      }

      index++;
      // we could break out of the while loop here, once the entire image has been streamed once
      if (index == numberOfStreamDivisions) index = 0;

      //otbAppLogINFO(""<<localNbSamples<<" samples added to the training set.");
    }

    otbAppLogINFO("The final training set contains "<<totalSamples<<" samples.");

    /*******************************************/
    /*           Learning                      */
    /*******************************************/
    otbAppLogINFO("-- LEARNING --");

    EstimatorType::Pointer estimator = EstimatorType::New();

    estimator->SetListSample(sampleList);
    EstimatorType::SizeType size;
    size[0]=GetParameterInt("sx");
    size[1]=GetParameterInt("sy");
    estimator->SetMapSize(size);
    EstimatorType::SizeType radius;
    radius[0] = GetParameterInt("nx");
    radius[1] = GetParameterInt("ny");
    estimator->SetNeighborhoodSizeInit(radius);
    estimator->SetNumberOfIterations(GetParameterInt("ni"));
    estimator->SetBetaInit(GetParameterFloat("bi"));
    estimator->SetBetaEnd(GetParameterFloat("bf"));
    estimator->SetMaxWeight(GetParameterFloat("iv"));

    AddProcess(estimator,"Learning");
    estimator->Update();

    m_SOMMap = estimator->GetOutput();
    SetParameterOutputImage<DoubleVectorImageType>("som",m_SOMMap);

    /*******************************************/
    /*           Classification                */
    /*******************************************/
    otbAppLogINFO("-- CLASSIFICATION --");

    m_Classifier = ClassificationFilterType::New();
    m_Classifier->SetInput(input);
    m_Classifier->SetInputMask(mask);
    m_Classifier->SetMap(m_SOMMap);

    AddProcess(m_Classifier,"Classification");

    SetParameterOutputImage<LabeledImageType>("out",m_Classifier->GetOutput());
    }

  SOMMapType::Pointer m_SOMMap;
  ClassificationFilterType::Pointer m_Classifier;
};

}
}

OTB_APPLICATION_EXPORT(otb::Wrapper::SOMClassification)

int main(int argc, char * argv[])
{

}
<|endoftext|>


<commit_before>//===- Disassembler.cpp - Disassembler for hex strings --------------------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// This class implements the disassembler of strings of bytes written in
// hexadecimal, from standard input or from a file.
//
//===----------------------------------------------------------------------===//

#include "Disassembler.h"

#include "llvm/ADT/OwningPtr.h"
#include "llvm/MC/MCAsmInfo.h"
#include "llvm/MC/MCDisassembler.h"
#include "llvm/MC/MCInst.h"
#include "llvm/MC/MCInstPrinter.h"
#include "llvm/Target/TargetRegistry.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/MemoryObject.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Support/SourceMgr.h"
using namespace llvm;

typedef std::vector<std::pair<unsigned char, const char*> > ByteArrayTy;

namespace {
class VectorMemoryObject : public MemoryObject {
private:
  const ByteArrayTy &Bytes;
public:
  VectorMemoryObject(const ByteArrayTy &bytes) : Bytes(bytes) {}

  uint64_t getBase() const { return 0; }
  uint64_t getExtent() const { return Bytes.size(); }

  int readByte(uint64_t Addr, uint8_t *Byte) const {
    if (Addr > getExtent())
      return -1;
    *Byte = Bytes[Addr].first;
    return 0;
  }
};
}

static bool PrintInst(const llvm::MCDisassembler &DisAsm,
                      llvm::MCInstPrinter &Printer, const ByteArrayTy &Bytes,
                      SourceMgr &SM) {
  // Wrap the vector in a MemoryObject.
  VectorMemoryObject memoryObject(Bytes);

  // Disassemble it to a string and get the size of the instruction.
  MCInst Inst;
  uint64_t Size;

  if (!DisAsm.getInstruction(Inst, Size, memoryObject, 0,
                             /*REMOVE*/ nulls())) {
    SM.PrintMessage(SMLoc::getFromPointer(Bytes[0].second),
                    "invalid instruction encoding", "error");
    return true;
  }

  Printer.printInst(&Inst);
  outs() << "\n";

  // If the disassembled instruction was smaller than the number of bytes we
  // read, reject the excess bytes.
  if (Bytes.size() != Size) {
    SM.PrintMessage(SMLoc::getFromPointer(Bytes[Size].second),
                    "excess data detected in input", "error");
    return true;
  }

  return false;
}

int Disassembler::disassemble(const Target &T, const std::string &Triple,
                                 MemoryBuffer &Buffer) {
  // Set up disassembler.
  llvm::OwningPtr<const llvm::MCAsmInfo> AsmInfo(T.createAsmInfo(Triple));

  if (!AsmInfo) {
    errs() << "error: no assembly info for target " << Triple << "\n";
    return -1;
  }

  llvm::OwningPtr<const llvm::MCDisassembler> DisAsm(T.createMCDisassembler());
  if (!DisAsm) {
    errs() << "error: no disassembler for target " << Triple << "\n";
    return -1;
  }

  llvm::MCInstPrinter *InstPrinter = T.createMCInstPrinter(0, *AsmInfo, outs());

  if (!InstPrinter) {
    errs() << "error: no instruction printer for target " << Triple << '\n';
    return -1;
  }

  bool ErrorOccurred = false;

  SourceMgr SM;
  SM.AddNewSourceBuffer(&Buffer, SMLoc());

  // Convert the input to a vector for disassembly.
  ByteArrayTy ByteArray;

  StringRef Str = Buffer.getBuffer();
  while (!Str.empty()) {
    // Strip horizontal whitespace.
    if (size_t Pos = Str.find_first_not_of(" \t\r")) {
      Str = Str.substr(Pos);
      continue;
    }

    // If this is the end of a line or start of a comment, process the
    // instruction we have so far.
    if (Str[0] == '\n' || Str[0] == '#') {
      // If we have bytes to process, do so.
      if (!ByteArray.empty()) {
        ErrorOccurred |= PrintInst(*DisAsm, *InstPrinter, ByteArray, SM);
        ByteArray.clear();
      }

      // Strip to the end of line if we already processed any bytes on this
      // line.  This strips the comment and/or the \n.
      if (Str[0] == '\n')
        Str = Str.substr(1);
      else {
        Str = Str.substr(Str.find_first_of('\n'));
        if (!Str.empty())
          Str = Str.substr(1);
      }
      continue;
    }

    // Get the current token.
    size_t Next = Str.find_first_of(" \t\n\r#");
    StringRef Value = Str.substr(0, Next);

    // Convert to a byte and add to the byte vector.
    unsigned ByteVal;
    if (Value.getAsInteger(0, ByteVal) || ByteVal > 255) {
      // If we have an error, print it and skip to the end of line.
      SM.PrintMessage(SMLoc::getFromPointer(Value.data()),
                                 "invalid input token", "error");
      ErrorOccurred = true;
      Str = Str.substr(Str.find('\n'));
      ByteArray.clear();
      continue;
    }

    ByteArray.push_back(std::make_pair((unsigned char)ByteVal, Value.data()));
    Str = Str.substr(Next);
  }

  if (!ByteArray.empty())
    ErrorOccurred |= PrintInst(*DisAsm, *InstPrinter, ByteArray, SM);

  return ErrorOccurred;
}
<commit_msg>Fixed the disassembler so it accepts multiple instructions on a single line.  Also made it a bit more forgiving when it reports errors.<commit_after>//===- Disassembler.cpp - Disassembler for hex strings --------------------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// This class implements the disassembler of strings of bytes written in
// hexadecimal, from standard input or from a file.
//
//===----------------------------------------------------------------------===//

#include "Disassembler.h"

#include "llvm/ADT/OwningPtr.h"
#include "llvm/MC/MCAsmInfo.h"
#include "llvm/MC/MCDisassembler.h"
#include "llvm/MC/MCInst.h"
#include "llvm/MC/MCInstPrinter.h"
#include "llvm/Target/TargetRegistry.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/MemoryObject.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Support/SourceMgr.h"
using namespace llvm;

typedef std::vector<std::pair<unsigned char, const char*> > ByteArrayTy;

namespace {
class VectorMemoryObject : public MemoryObject {
private:
  const ByteArrayTy &Bytes;
public:
  VectorMemoryObject(const ByteArrayTy &bytes) : Bytes(bytes) {}

  uint64_t getBase() const { return 0; }
  uint64_t getExtent() const { return Bytes.size(); }

  int readByte(uint64_t Addr, uint8_t *Byte) const {
    if (Addr > getExtent())
      return -1;
    *Byte = Bytes[Addr].first;
    return 0;
  }
};
}

static bool PrintInsts(const llvm::MCDisassembler &DisAsm,
                      llvm::MCInstPrinter &Printer, const ByteArrayTy &Bytes,
                      SourceMgr &SM) {
  // Wrap the vector in a MemoryObject.
  VectorMemoryObject memoryObject(Bytes);

  // Disassemble it to strings.
  uint64_t Size;
  uint64_t Index;

  for (Index = 0; Index < Bytes.size(); Index += Size) {
    MCInst Inst;

    if (DisAsm.getInstruction(Inst, Size, memoryObject, Index,
                               /*REMOVE*/ nulls())) {
      Printer.printInst(&Inst);
      outs() << "\n";
    }
    else {
      SM.PrintMessage(SMLoc::getFromPointer(Bytes[Index].second),
                      "invalid instruction encoding", "warning");
      if (Size == 0)
        Size = 1; // skip illegible bytes
    }
  }

  return false;
}

int Disassembler::disassemble(const Target &T, const std::string &Triple,
                                 MemoryBuffer &Buffer) {
  // Set up disassembler.
  llvm::OwningPtr<const llvm::MCAsmInfo> AsmInfo(T.createAsmInfo(Triple));

  if (!AsmInfo) {
    errs() << "error: no assembly info for target " << Triple << "\n";
    return -1;
  }

  llvm::OwningPtr<const llvm::MCDisassembler> DisAsm(T.createMCDisassembler());
  if (!DisAsm) {
    errs() << "error: no disassembler for target " << Triple << "\n";
    return -1;
  }

  llvm::MCInstPrinter *InstPrinter = T.createMCInstPrinter(0, *AsmInfo, outs());

  if (!InstPrinter) {
    errs() << "error: no instruction printer for target " << Triple << '\n';
    return -1;
  }

  bool ErrorOccurred = false;

  SourceMgr SM;
  SM.AddNewSourceBuffer(&Buffer, SMLoc());

  // Convert the input to a vector for disassembly.
  ByteArrayTy ByteArray;

  StringRef Str = Buffer.getBuffer();
  while (!Str.empty()) {
    // Strip horizontal whitespace.
    if (size_t Pos = Str.find_first_not_of(" \t\r")) {
      Str = Str.substr(Pos);
      continue;
    }

    // If this is the end of a line or start of a comment, remove the rest of
    // the line.
    if (Str[0] == '\n' || Str[0] == '#') {
      // Strip to the end of line if we already processed any bytes on this
      // line.  This strips the comment and/or the \n.
      if (Str[0] == '\n')
        Str = Str.substr(1);
      else {
        Str = Str.substr(Str.find_first_of('\n'));
        if (!Str.empty())
          Str = Str.substr(1);
      }
      continue;
    }

    // Get the current token.
    size_t Next = Str.find_first_of(" \t\n\r#");
    StringRef Value = Str.substr(0, Next);

    // Convert to a byte and add to the byte vector.
    unsigned ByteVal;
    if (Value.getAsInteger(0, ByteVal) || ByteVal > 255) {
      // If we have an error, print it and skip to the end of line.
      SM.PrintMessage(SMLoc::getFromPointer(Value.data()),
                                 "invalid input token", "error");
      ErrorOccurred = true;
      Str = Str.substr(Str.find('\n'));
      ByteArray.clear();
      continue;
    }

    ByteArray.push_back(std::make_pair((unsigned char)ByteVal, Value.data()));
    Str = Str.substr(Next);
  }

  if (!ByteArray.empty())
    ErrorOccurred |= PrintInsts(*DisAsm, *InstPrinter, ByteArray, SM);

  return ErrorOccurred;
}
<|endoftext|>


<commit_before>//===-LTOCodeGenerator.cpp - LLVM Link Time Optimizer ---------------------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// This file implements the Link Time Optimization library. This library is
// intended to be used by linker to optimize code at link time.
//
//===----------------------------------------------------------------------===//

#include "LTOModule.h"
#include "LTOCodeGenerator.h"

#include "llvm/Constants.h"
#include "llvm/DerivedTypes.h"
#include "llvm/Linker.h"
#include "llvm/LLVMContext.h"
#include "llvm/Module.h"
#include "llvm/PassManager.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/ADT/Triple.h"
#include "llvm/Analysis/Passes.h"
#include "llvm/Bitcode/ReaderWriter.h"
#include "llvm/MC/MCAsmInfo.h"
#include "llvm/MC/MCContext.h"
#include "llvm/Target/Mangler.h"
#include "llvm/Target/SubtargetFeature.h"
#include "llvm/Target/TargetOptions.h"
#include "llvm/Target/TargetData.h"
#include "llvm/Target/TargetMachine.h"
#include "llvm/Target/TargetRegistry.h"
#include "llvm/Target/TargetSelect.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/FormattedStream.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/StandardPasses.h"
#include "llvm/Support/SystemUtils.h"
#include "llvm/Support/ToolOutputFile.h"
#include "llvm/Support/Host.h"
#include "llvm/Support/Program.h"
#include "llvm/Support/Signals.h"
#include "llvm/Config/config.h"
#include <cstdlib>
#include <unistd.h>
#include <fcntl.h>

using namespace llvm;

static cl::opt<bool> DisableInline("disable-inlining",
  cl::desc("Do not run the inliner pass"));

const char* LTOCodeGenerator::getVersionString()
{
#ifdef LLVM_VERSION_INFO
    return PACKAGE_NAME " version " PACKAGE_VERSION ", " LLVM_VERSION_INFO;
#else
    return PACKAGE_NAME " version " PACKAGE_VERSION;
#endif
}

LTOCodeGenerator::LTOCodeGenerator()
    : _context(getGlobalContext()),
      _linker("LinkTimeOptimizer", "ld-temp.o", _context), _target(NULL),
      _emitDwarfDebugInfo(false), _scopeRestrictionsDone(false),
      _codeModel(LTO_CODEGEN_PIC_MODEL_DYNAMIC),
      _nativeObjectFile(NULL), _assemblerPath(NULL)
{
    InitializeAllTargets();
    InitializeAllAsmPrinters();
}

LTOCodeGenerator::~LTOCodeGenerator()
{
    delete _target;
    delete _nativeObjectFile;
}

bool LTOCodeGenerator::addModule(LTOModule* mod, std::string& errMsg)
{
    return _linker.LinkInModule(mod->getLLVVMModule(), &errMsg);
}

bool LTOCodeGenerator::setDebugInfo(lto_debug_model debug, std::string& errMsg)
{
    switch (debug) {
        case LTO_DEBUG_MODEL_NONE:
            _emitDwarfDebugInfo = false;
            return false;

        case LTO_DEBUG_MODEL_DWARF:
            _emitDwarfDebugInfo = true;
            return false;
    }
    errMsg = "unknown debug format";
    return true;
}

bool LTOCodeGenerator::setCodePICModel(lto_codegen_model model,
                                       std::string& errMsg)
{
    switch (model) {
        case LTO_CODEGEN_PIC_MODEL_STATIC:
        case LTO_CODEGEN_PIC_MODEL_DYNAMIC:
        case LTO_CODEGEN_PIC_MODEL_DYNAMIC_NO_PIC:
            _codeModel = model;
            return false;
    }
    errMsg = "unknown pic model";
    return true;
}

void LTOCodeGenerator::setCpu(const char* mCpu)
{
  _mCpu = mCpu;
}

void LTOCodeGenerator::setAssemblerPath(const char* path)
{
    if ( _assemblerPath )
        delete _assemblerPath;
    _assemblerPath = new sys::Path(path);
}

void LTOCodeGenerator::setAssemblerArgs(const char** args, int nargs)
{
  for (int i = 0; i < nargs; ++i) {
    const char *arg = args[i];
    _assemblerArgs.push_back(arg);
  }
}

void LTOCodeGenerator::addMustPreserveSymbol(const char* sym)
{
    _mustPreserveSymbols[sym] = 1;
}

bool LTOCodeGenerator::writeMergedModules(const char *path,
                                          std::string &errMsg) {
  if (determineTarget(errMsg))
    return true;

  // mark which symbols can not be internalized
  applyScopeRestrictions();

  // create output file
  std::string ErrInfo;
  tool_output_file Out(path, ErrInfo,
                       raw_fd_ostream::F_Binary);
  if (!ErrInfo.empty()) {
    errMsg = "could not open bitcode file for writing: ";
    errMsg += path;
    return true;
  }

  // write bitcode to it
  WriteBitcodeToFile(_linker.getModule(), Out.os());
  Out.os().close();

  if (Out.os().has_error()) {
    errMsg = "could not write bitcode file: ";
    errMsg += path;
    Out.os().clear_error();
    return true;
  }

  Out.keep();
  return false;
}

const void* LTOCodeGenerator::compile(size_t* length, std::string& errMsg)
{
    // make unique temp .s file to put generated assembly code
    sys::Path uniqueAsmPath("lto-llvm.s");
    if ( uniqueAsmPath.createTemporaryFileOnDisk(true, &errMsg) )
        return NULL;
    sys::RemoveFileOnSignal(uniqueAsmPath);

    // generate assembly code
    bool genResult = false;
    {
      tool_output_file asmFile(uniqueAsmPath.c_str(), errMsg);
      if (!errMsg.empty())
        return NULL;
      genResult = this->generateAssemblyCode(asmFile.os(), errMsg);
      asmFile.os().close();
      if (asmFile.os().has_error()) {
        asmFile.os().clear_error();
        return NULL;
      }
      asmFile.keep();
    }
    if ( genResult ) {
        uniqueAsmPath.eraseFromDisk();
        return NULL;
    }

    // make unique temp .o file to put generated object file
    sys::PathWithStatus uniqueObjPath("lto-llvm.o");
    if ( uniqueObjPath.createTemporaryFileOnDisk(false, &errMsg) ) {
        uniqueAsmPath.eraseFromDisk();
        return NULL;
    }
    sys::RemoveFileOnSignal(uniqueObjPath);

    // assemble the assembly code
    const std::string& uniqueObjStr = uniqueObjPath.str();
    bool asmResult = this->assemble(uniqueAsmPath.str(), uniqueObjStr, errMsg);
    if ( !asmResult ) {
        // remove old buffer if compile() called twice
        delete _nativeObjectFile;

        // read .o file into memory buffer
        _nativeObjectFile = MemoryBuffer::getFile(uniqueObjStr.c_str(),&errMsg);
    }

    // remove temp files
    uniqueAsmPath.eraseFromDisk();
    uniqueObjPath.eraseFromDisk();

    // return buffer, unless error
    if ( _nativeObjectFile == NULL )
        return NULL;
    *length = _nativeObjectFile->getBufferSize();
    return _nativeObjectFile->getBufferStart();
}

bool LTOCodeGenerator::assemble(const std::string& asmPath,
                                const std::string& objPath, std::string& errMsg)
{
    sys::Path tool;
    bool needsCompilerOptions = true;
    if ( _assemblerPath ) {
        tool = *_assemblerPath;
        needsCompilerOptions = false;
    } else {
        // find compiler driver
        tool = sys::Program::FindProgramByName("gcc");
        if ( tool.isEmpty() ) {
            errMsg = "can't locate gcc";
            return true;
        }
    }

    // build argument list
    std::vector<const char*> args;
    llvm::Triple targetTriple(_linker.getModule()->getTargetTriple());
    const char *arch = targetTriple.getArchNameForAssembler();

    args.push_back(tool.c_str());

    if (targetTriple.getOS() == Triple::Darwin) {
        // darwin specific command line options
        if (arch != NULL) {
            args.push_back("-arch");
            args.push_back(arch);
        }
        // add -static to assembler command line when code model requires
        if ( (_assemblerPath != NULL) &&
            (_codeModel == LTO_CODEGEN_PIC_MODEL_STATIC) )
            args.push_back("-static");
    }
    if ( needsCompilerOptions ) {
        args.push_back("-c");
        args.push_back("-x");
        args.push_back("assembler");
    } else {
        for (std::vector<std::string>::iterator I = _assemblerArgs.begin(),
               E = _assemblerArgs.end(); I != E; ++I) {
            args.push_back(I->c_str());
        }
    }
    args.push_back("-o");
    args.push_back(objPath.c_str());
    args.push_back(asmPath.c_str());
    args.push_back(0);

    // invoke assembler
    if ( sys::Program::ExecuteAndWait(tool, &args[0], 0, 0, 0, 0, &errMsg) ) {
        errMsg = "error in assembly";
        return true;
    }
    return false; // success
}

bool LTOCodeGenerator::determineTarget(std::string& errMsg)
{
    if ( _target == NULL ) {
        std::string Triple = _linker.getModule()->getTargetTriple();
        if (Triple.empty())
          Triple = sys::getHostTriple();

        // create target machine from info for merged modules
        const Target *march = TargetRegistry::lookupTarget(Triple, errMsg);
        if ( march == NULL )
            return true;

        // The relocation model is actually a static member of TargetMachine
        // and needs to be set before the TargetMachine is instantiated.
        switch( _codeModel ) {
        case LTO_CODEGEN_PIC_MODEL_STATIC:
            TargetMachine::setRelocationModel(Reloc::Static);
            break;
        case LTO_CODEGEN_PIC_MODEL_DYNAMIC:
            TargetMachine::setRelocationModel(Reloc::PIC_);
            break;
        case LTO_CODEGEN_PIC_MODEL_DYNAMIC_NO_PIC:
            TargetMachine::setRelocationModel(Reloc::DynamicNoPIC);
            break;
        }

        // construct LTModule, hand over ownership of module and target
        SubtargetFeatures Features;
        Features.getDefaultSubtargetFeatures(_mCpu, llvm::Triple(Triple));
        std::string FeatureStr = Features.getString();
        _target = march->createTargetMachine(Triple, FeatureStr);
    }
    return false;
}

void LTOCodeGenerator::applyScopeRestrictions() {
  if (_scopeRestrictionsDone) return;
  Module *mergedModule = _linker.getModule();

  // Start off with a verification pass.
  PassManager passes;
  passes.add(createVerifierPass());

  // mark which symbols can not be internalized
  if (!_mustPreserveSymbols.empty()) {
    MCContext Context(*_target->getMCAsmInfo());
    Mangler mangler(Context, *_target->getTargetData());
    std::vector<const char*> mustPreserveList;
    for (Module::iterator f = mergedModule->begin(),
         e = mergedModule->end(); f != e; ++f) {
      if (!f->isDeclaration() &&
          _mustPreserveSymbols.count(mangler.getNameWithPrefix(f)))
        mustPreserveList.push_back(::strdup(f->getNameStr().c_str()));
    }
    for (Module::global_iterator v = mergedModule->global_begin(),
         e = mergedModule->global_end(); v !=  e; ++v) {
      if (!v->isDeclaration() &&
          _mustPreserveSymbols.count(mangler.getNameWithPrefix(v)))
        mustPreserveList.push_back(::strdup(v->getNameStr().c_str()));
    }
    passes.add(createInternalizePass(mustPreserveList));
  }

  // apply scope restrictions
  passes.run(*mergedModule);

  _scopeRestrictionsDone = true;
}

/// Optimize merged modules using various IPO passes
bool LTOCodeGenerator::generateAssemblyCode(raw_ostream& out,
                                            std::string& errMsg)
{
    if ( this->determineTarget(errMsg) )
        return true;

    // mark which symbols can not be internalized
    this->applyScopeRestrictions();

    Module* mergedModule = _linker.getModule();

    // if options were requested, set them
    if ( !_codegenOptions.empty() )
        cl::ParseCommandLineOptions(_codegenOptions.size(),
                                    const_cast<char **>(&_codegenOptions[0]));

    // Instantiate the pass manager to organize the passes.
    PassManager passes;

    // Start off with a verification pass.
    passes.add(createVerifierPass());

    // Add an appropriate TargetData instance for this module...
    passes.add(new TargetData(*_target->getTargetData()));

    createStandardLTOPasses(&passes, /*Internalize=*/ false, !DisableInline,
                            /*VerifyEach=*/ false);

    // Make sure everything is still good.
    passes.add(createVerifierPass());

    FunctionPassManager* codeGenPasses = new FunctionPassManager(mergedModule);

    codeGenPasses->add(new TargetData(*_target->getTargetData()));

    formatted_raw_ostream Out(out);

    if (_target->addPassesToEmitFile(*codeGenPasses, Out,
                                     TargetMachine::CGFT_AssemblyFile,
                                     CodeGenOpt::Aggressive)) {
      errMsg = "target file type not supported";
      return true;
    }

    // Run our queue of passes all at once now, efficiently.
    passes.run(*mergedModule);

    // Run the code generator, and write assembly file
    codeGenPasses->doInitialization();

    for (Module::iterator
           it = mergedModule->begin(), e = mergedModule->end(); it != e; ++it)
      if (!it->isDeclaration())
        codeGenPasses->run(*it);

    codeGenPasses->doFinalization();

    return false; // success
}

/// Optimize merged modules using various IPO passes
void LTOCodeGenerator::setCodeGenDebugOptions(const char* options)
{
    for (std::pair<StringRef, StringRef> o = getToken(options);
         !o.first.empty(); o = getToken(o.second)) {
        // ParseCommandLineOptions() expects argv[0] to be program name.
        // Lazily add that.
        if ( _codegenOptions.empty() )
            _codegenOptions.push_back("libLTO");
        _codegenOptions.push_back(strdup(o.first.str().c_str()));
    }
}
<commit_msg>Do not try luck by using given name to create temporary file. In parallel builds it may not work.<commit_after>//===-LTOCodeGenerator.cpp - LLVM Link Time Optimizer ---------------------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// This file implements the Link Time Optimization library. This library is
// intended to be used by linker to optimize code at link time.
//
//===----------------------------------------------------------------------===//

#include "LTOModule.h"
#include "LTOCodeGenerator.h"

#include "llvm/Constants.h"
#include "llvm/DerivedTypes.h"
#include "llvm/Linker.h"
#include "llvm/LLVMContext.h"
#include "llvm/Module.h"
#include "llvm/PassManager.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/ADT/Triple.h"
#include "llvm/Analysis/Passes.h"
#include "llvm/Bitcode/ReaderWriter.h"
#include "llvm/MC/MCAsmInfo.h"
#include "llvm/MC/MCContext.h"
#include "llvm/Target/Mangler.h"
#include "llvm/Target/SubtargetFeature.h"
#include "llvm/Target/TargetOptions.h"
#include "llvm/Target/TargetData.h"
#include "llvm/Target/TargetMachine.h"
#include "llvm/Target/TargetRegistry.h"
#include "llvm/Target/TargetSelect.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/FormattedStream.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/StandardPasses.h"
#include "llvm/Support/SystemUtils.h"
#include "llvm/Support/ToolOutputFile.h"
#include "llvm/Support/Host.h"
#include "llvm/Support/Program.h"
#include "llvm/Support/Signals.h"
#include "llvm/Config/config.h"
#include <cstdlib>
#include <unistd.h>
#include <fcntl.h>

using namespace llvm;

static cl::opt<bool> DisableInline("disable-inlining",
  cl::desc("Do not run the inliner pass"));

const char* LTOCodeGenerator::getVersionString()
{
#ifdef LLVM_VERSION_INFO
    return PACKAGE_NAME " version " PACKAGE_VERSION ", " LLVM_VERSION_INFO;
#else
    return PACKAGE_NAME " version " PACKAGE_VERSION;
#endif
}

LTOCodeGenerator::LTOCodeGenerator()
    : _context(getGlobalContext()),
      _linker("LinkTimeOptimizer", "ld-temp.o", _context), _target(NULL),
      _emitDwarfDebugInfo(false), _scopeRestrictionsDone(false),
      _codeModel(LTO_CODEGEN_PIC_MODEL_DYNAMIC),
      _nativeObjectFile(NULL), _assemblerPath(NULL)
{
    InitializeAllTargets();
    InitializeAllAsmPrinters();
}

LTOCodeGenerator::~LTOCodeGenerator()
{
    delete _target;
    delete _nativeObjectFile;
}

bool LTOCodeGenerator::addModule(LTOModule* mod, std::string& errMsg)
{
    return _linker.LinkInModule(mod->getLLVVMModule(), &errMsg);
}

bool LTOCodeGenerator::setDebugInfo(lto_debug_model debug, std::string& errMsg)
{
    switch (debug) {
        case LTO_DEBUG_MODEL_NONE:
            _emitDwarfDebugInfo = false;
            return false;

        case LTO_DEBUG_MODEL_DWARF:
            _emitDwarfDebugInfo = true;
            return false;
    }
    errMsg = "unknown debug format";
    return true;
}

bool LTOCodeGenerator::setCodePICModel(lto_codegen_model model,
                                       std::string& errMsg)
{
    switch (model) {
        case LTO_CODEGEN_PIC_MODEL_STATIC:
        case LTO_CODEGEN_PIC_MODEL_DYNAMIC:
        case LTO_CODEGEN_PIC_MODEL_DYNAMIC_NO_PIC:
            _codeModel = model;
            return false;
    }
    errMsg = "unknown pic model";
    return true;
}

void LTOCodeGenerator::setCpu(const char* mCpu)
{
  _mCpu = mCpu;
}

void LTOCodeGenerator::setAssemblerPath(const char* path)
{
    if ( _assemblerPath )
        delete _assemblerPath;
    _assemblerPath = new sys::Path(path);
}

void LTOCodeGenerator::setAssemblerArgs(const char** args, int nargs)
{
  for (int i = 0; i < nargs; ++i) {
    const char *arg = args[i];
    _assemblerArgs.push_back(arg);
  }
}

void LTOCodeGenerator::addMustPreserveSymbol(const char* sym)
{
    _mustPreserveSymbols[sym] = 1;
}

bool LTOCodeGenerator::writeMergedModules(const char *path,
                                          std::string &errMsg) {
  if (determineTarget(errMsg))
    return true;

  // mark which symbols can not be internalized
  applyScopeRestrictions();

  // create output file
  std::string ErrInfo;
  tool_output_file Out(path, ErrInfo,
                       raw_fd_ostream::F_Binary);
  if (!ErrInfo.empty()) {
    errMsg = "could not open bitcode file for writing: ";
    errMsg += path;
    return true;
  }

  // write bitcode to it
  WriteBitcodeToFile(_linker.getModule(), Out.os());
  Out.os().close();

  if (Out.os().has_error()) {
    errMsg = "could not write bitcode file: ";
    errMsg += path;
    Out.os().clear_error();
    return true;
  }

  Out.keep();
  return false;
}

const void* LTOCodeGenerator::compile(size_t* length, std::string& errMsg)
{
    // make unique temp .s file to put generated assembly code
    sys::Path uniqueAsmPath("lto-llvm.s");
    if ( uniqueAsmPath.createTemporaryFileOnDisk(false, &errMsg) )
        return NULL;
    sys::RemoveFileOnSignal(uniqueAsmPath);

    // generate assembly code
    bool genResult = false;
    {
      tool_output_file asmFile(uniqueAsmPath.c_str(), errMsg);
      if (!errMsg.empty())
        return NULL;
      genResult = this->generateAssemblyCode(asmFile.os(), errMsg);
      asmFile.os().close();
      if (asmFile.os().has_error()) {
        asmFile.os().clear_error();
        return NULL;
      }
      asmFile.keep();
    }
    if ( genResult ) {
        uniqueAsmPath.eraseFromDisk();
        return NULL;
    }

    // make unique temp .o file to put generated object file
    sys::PathWithStatus uniqueObjPath("lto-llvm.o");
    if ( uniqueObjPath.createTemporaryFileOnDisk(false, &errMsg) ) {
        uniqueAsmPath.eraseFromDisk();
        return NULL;
    }
    sys::RemoveFileOnSignal(uniqueObjPath);

    // assemble the assembly code
    const std::string& uniqueObjStr = uniqueObjPath.str();
    bool asmResult = this->assemble(uniqueAsmPath.str(), uniqueObjStr, errMsg);
    if ( !asmResult ) {
        // remove old buffer if compile() called twice
        delete _nativeObjectFile;

        // read .o file into memory buffer
        _nativeObjectFile = MemoryBuffer::getFile(uniqueObjStr.c_str(),&errMsg);
    }

    // remove temp files
    uniqueAsmPath.eraseFromDisk();
    uniqueObjPath.eraseFromDisk();

    // return buffer, unless error
    if ( _nativeObjectFile == NULL )
        return NULL;
    *length = _nativeObjectFile->getBufferSize();
    return _nativeObjectFile->getBufferStart();
}

bool LTOCodeGenerator::assemble(const std::string& asmPath,
                                const std::string& objPath, std::string& errMsg)
{
    sys::Path tool;
    bool needsCompilerOptions = true;
    if ( _assemblerPath ) {
        tool = *_assemblerPath;
        needsCompilerOptions = false;
    } else {
        // find compiler driver
        tool = sys::Program::FindProgramByName("gcc");
        if ( tool.isEmpty() ) {
            errMsg = "can't locate gcc";
            return true;
        }
    }

    // build argument list
    std::vector<const char*> args;
    llvm::Triple targetTriple(_linker.getModule()->getTargetTriple());
    const char *arch = targetTriple.getArchNameForAssembler();

    args.push_back(tool.c_str());

    if (targetTriple.getOS() == Triple::Darwin) {
        // darwin specific command line options
        if (arch != NULL) {
            args.push_back("-arch");
            args.push_back(arch);
        }
        // add -static to assembler command line when code model requires
        if ( (_assemblerPath != NULL) &&
            (_codeModel == LTO_CODEGEN_PIC_MODEL_STATIC) )
            args.push_back("-static");
    }
    if ( needsCompilerOptions ) {
        args.push_back("-c");
        args.push_back("-x");
        args.push_back("assembler");
    } else {
        for (std::vector<std::string>::iterator I = _assemblerArgs.begin(),
               E = _assemblerArgs.end(); I != E; ++I) {
            args.push_back(I->c_str());
        }
    }
    args.push_back("-o");
    args.push_back(objPath.c_str());
    args.push_back(asmPath.c_str());
    args.push_back(0);

    // invoke assembler
    if ( sys::Program::ExecuteAndWait(tool, &args[0], 0, 0, 0, 0, &errMsg) ) {
        errMsg = "error in assembly";
        return true;
    }
    return false; // success
}

bool LTOCodeGenerator::determineTarget(std::string& errMsg)
{
    if ( _target == NULL ) {
        std::string Triple = _linker.getModule()->getTargetTriple();
        if (Triple.empty())
          Triple = sys::getHostTriple();

        // create target machine from info for merged modules
        const Target *march = TargetRegistry::lookupTarget(Triple, errMsg);
        if ( march == NULL )
            return true;

        // The relocation model is actually a static member of TargetMachine
        // and needs to be set before the TargetMachine is instantiated.
        switch( _codeModel ) {
        case LTO_CODEGEN_PIC_MODEL_STATIC:
            TargetMachine::setRelocationModel(Reloc::Static);
            break;
        case LTO_CODEGEN_PIC_MODEL_DYNAMIC:
            TargetMachine::setRelocationModel(Reloc::PIC_);
            break;
        case LTO_CODEGEN_PIC_MODEL_DYNAMIC_NO_PIC:
            TargetMachine::setRelocationModel(Reloc::DynamicNoPIC);
            break;
        }

        // construct LTModule, hand over ownership of module and target
        SubtargetFeatures Features;
        Features.getDefaultSubtargetFeatures(_mCpu, llvm::Triple(Triple));
        std::string FeatureStr = Features.getString();
        _target = march->createTargetMachine(Triple, FeatureStr);
    }
    return false;
}

void LTOCodeGenerator::applyScopeRestrictions() {
  if (_scopeRestrictionsDone) return;
  Module *mergedModule = _linker.getModule();

  // Start off with a verification pass.
  PassManager passes;
  passes.add(createVerifierPass());

  // mark which symbols can not be internalized
  if (!_mustPreserveSymbols.empty()) {
    MCContext Context(*_target->getMCAsmInfo());
    Mangler mangler(Context, *_target->getTargetData());
    std::vector<const char*> mustPreserveList;
    for (Module::iterator f = mergedModule->begin(),
         e = mergedModule->end(); f != e; ++f) {
      if (!f->isDeclaration() &&
          _mustPreserveSymbols.count(mangler.getNameWithPrefix(f)))
        mustPreserveList.push_back(::strdup(f->getNameStr().c_str()));
    }
    for (Module::global_iterator v = mergedModule->global_begin(),
         e = mergedModule->global_end(); v !=  e; ++v) {
      if (!v->isDeclaration() &&
          _mustPreserveSymbols.count(mangler.getNameWithPrefix(v)))
        mustPreserveList.push_back(::strdup(v->getNameStr().c_str()));
    }
    passes.add(createInternalizePass(mustPreserveList));
  }

  // apply scope restrictions
  passes.run(*mergedModule);

  _scopeRestrictionsDone = true;
}

/// Optimize merged modules using various IPO passes
bool LTOCodeGenerator::generateAssemblyCode(raw_ostream& out,
                                            std::string& errMsg)
{
    if ( this->determineTarget(errMsg) )
        return true;

    // mark which symbols can not be internalized
    this->applyScopeRestrictions();

    Module* mergedModule = _linker.getModule();

    // if options were requested, set them
    if ( !_codegenOptions.empty() )
        cl::ParseCommandLineOptions(_codegenOptions.size(),
                                    const_cast<char **>(&_codegenOptions[0]));

    // Instantiate the pass manager to organize the passes.
    PassManager passes;

    // Start off with a verification pass.
    passes.add(createVerifierPass());

    // Add an appropriate TargetData instance for this module...
    passes.add(new TargetData(*_target->getTargetData()));

    createStandardLTOPasses(&passes, /*Internalize=*/ false, !DisableInline,
                            /*VerifyEach=*/ false);

    // Make sure everything is still good.
    passes.add(createVerifierPass());

    FunctionPassManager* codeGenPasses = new FunctionPassManager(mergedModule);

    codeGenPasses->add(new TargetData(*_target->getTargetData()));

    formatted_raw_ostream Out(out);

    if (_target->addPassesToEmitFile(*codeGenPasses, Out,
                                     TargetMachine::CGFT_AssemblyFile,
                                     CodeGenOpt::Aggressive)) {
      errMsg = "target file type not supported";
      return true;
    }

    // Run our queue of passes all at once now, efficiently.
    passes.run(*mergedModule);

    // Run the code generator, and write assembly file
    codeGenPasses->doInitialization();

    for (Module::iterator
           it = mergedModule->begin(), e = mergedModule->end(); it != e; ++it)
      if (!it->isDeclaration())
        codeGenPasses->run(*it);

    codeGenPasses->doFinalization();

    return false; // success
}

/// Optimize merged modules using various IPO passes
void LTOCodeGenerator::setCodeGenDebugOptions(const char* options)
{
    for (std::pair<StringRef, StringRef> o = getToken(options);
         !o.first.empty(); o = getToken(o.second)) {
        // ParseCommandLineOptions() expects argv[0] to be program name.
        // Lazily add that.
        if ( _codegenOptions.empty() )
            _codegenOptions.push_back("libLTO");
        _codegenOptions.push_back(strdup(o.first.str().c_str()));
    }
}
<|endoftext|>