mosesdecoder/moses/TranslationModel/CompactPT/MmapAllocator.h

// $Id$                                                                                                                               
// vim:tabstop=2                                                                                                                      
/***********************************************************************                                                              
Moses - factored phrase-based language decoder                                                                                        
Copyright (C) 2006 University of Edinburgh                                                                                            
                                                                                                                                      
This library is free software; you can redistribute it and/or                                                                         
modify it under the terms of the GNU Lesser General Public                                                                            
License as published by the Free Software Foundation; either                                                                          
version 2.1 of the License, or (at your option) any later version.                                                                    
                                                                                                                                      
This library is distributed in the hope that it will be useful,                                                                       
but WITHOUT ANY WARRANTY; without even the implied warranty of                                                                        
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU                                                                     
Lesser General Public License for more details.                                                                                       
                                                                                                                                      
You should have received a copy of the GNU Lesser General Public                                                                      
License along with this library; if not, write to the Free Software                                                                   
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA                                                        
***********************************************************************/  

#ifndef moses_MmapAllocator_h
#define moses_MmapAllocator_h

#include <limits>
#include <iostream>
#include <sys/mman.h>
#include <cstdio>
#include <unistd.h>

namespace Moses
{
    template <class T>
    class MmapAllocator
    {
      protected:    
        std::FILE* m_file_ptr;
        size_t m_file_desc;
        
        size_t m_page_size;
        size_t m_map_size;
        
        char* m_data_ptr;
        size_t m_data_offset;
        bool m_fixed;
        size_t* m_count;
        
      public:
        typedef T        value_type;
        typedef T*       pointer;
        typedef const T* const_pointer;
        typedef T&       reference;
        typedef const T& const_reference;
        typedef std::size_t    size_type;
        typedef std::ptrdiff_t difference_type;
    
        MmapAllocator() throw()
         : m_file_ptr(std::tmpfile()), m_file_desc(fileno(m_file_ptr)),
           m_page_size(sysconf(_SC_PAGE_SIZE)), m_map_size(0), m_data_ptr(0),
           m_data_offset(0), m_fixed(false), m_count(new size_t(0))
        { }
        
        MmapAllocator(std::FILE* f_ptr) throw()
         : m_file_ptr(f_ptr), m_file_desc(fileno(m_file_ptr)),
           m_page_size(sysconf(_SC_PAGE_SIZE)), m_map_size(0), m_data_ptr(0),
           m_data_offset(0), m_fixed(false), m_count(new size_t(0))
        { }
        
        MmapAllocator(std::FILE* f_ptr, size_t data_offset) throw()
         : m_file_ptr(f_ptr), m_file_desc(fileno(m_file_ptr)),
           m_page_size(sysconf(_SC_PAGE_SIZE)), m_map_size(0), m_data_ptr(0),
           m_data_offset(data_offset), m_fixed(true), m_count(new size_t(0))
        { }
        
        MmapAllocator(std::string fileName) throw()
         : m_file_ptr(std::fopen(fileName.c_str(), "wb+")), m_file_desc(fileno(m_file_ptr)),
           m_page_size(sysconf(_SC_PAGE_SIZE)), m_map_size(0), m_data_ptr(0),
           m_data_offset(0), m_fixed(false), m_count(new size_t(0))
        { }
            
        MmapAllocator(const MmapAllocator& c) throw()
         : m_file_ptr(c.m_file_ptr), m_file_desc(c.m_file_desc),
           m_page_size(c.m_page_size), m_map_size(c.m_map_size),
           m_data_ptr(c.m_data_ptr), m_data_offset(c.m_data_offset),
           m_fixed(c.m_fixed), m_count(c.m_count)
        { 
          (*m_count)++;
        }
        
        ~MmapAllocator() throw()
        {
            if(m_data_ptr && *m_count == 0)
            {
                munmap(m_data_ptr, m_map_size);
                if(!m_fixed && std::ftell(m_file_ptr) != -1)
                    std::fclose(m_file_ptr);
            }
            (*m_count)--;
        }
     
        template <class U>
        struct rebind {
            typedef MmapAllocator<U> other;
        };
     
        pointer address (reference value) const
        {
            return &value;
        }
     
        const_pointer address (const_reference value) const
        {
            return &value;
        }
     
        size_type max_size () const throw()
        {
            return std::numeric_limits<size_t>::max() / sizeof(value_type);
        }
     
        pointer allocate (size_type num, const void* = 0)
        {
            m_map_size = num * sizeof(T);
            
            if(!m_fixed)
            {
                size_t read = 0;
                read += ftruncate(m_file_desc, m_map_size);
                m_data_ptr = (char*)mmap(0, m_map_size, PROT_READ|PROT_WRITE, MAP_SHARED,
                                       m_file_desc, 0);
                if(m_data_ptr == MAP_FAILED)
                  std::cerr << "Error: mmapping" << std::endl;
                return (pointer)m_data_ptr;
            }
            else
            {
                size_t map_offset = (m_data_offset / m_page_size) * m_page_size;
                size_t relative_offset = m_data_offset - map_offset;
                
                size_t map_size = m_map_size + relative_offset;
            
                m_data_ptr = (char*)mmap(0, map_size, PROT_READ, MAP_SHARED,
                                       m_file_desc, map_offset);

                return (pointer)(m_data_ptr + relative_offset);
            }
        }
     
        void deallocate (pointer p, size_type num)
        {
            if(!m_fixed) {
                munmap(p, num * sizeof(T));
            }
            else {
                size_t map_offset = (m_data_offset / m_page_size) * m_page_size;
                size_t relative_offset = m_data_offset - map_offset;
                munmap((pointer)((char*)p - relative_offset), num * sizeof(T));    
            }
            
        }
    
        void construct (pointer p, const T& value)
        {
            if(!m_fixed)
                new(p) value_type(value);
        }
        void destroy (pointer p)
        {
            if(!m_fixed)
                p->~T();
        }
         
        template <class T1, class T2>
        friend bool operator== (const MmapAllocator<T1>&, const MmapAllocator<T2>&) throw();
    
        template <class T1, class T2>
        friend bool operator!= (const MmapAllocator<T1>&, const MmapAllocator<T2>&) throw();
    };
    
    template <class T1, class T2>
    bool operator== (const MmapAllocator<T1>& a1,
                     const MmapAllocator<T2>& a2) throw()
    {    
        bool equal = true;
        equal &= a1.m_file_ptr == a2.m_file_ptr;
        equal &= a1.m_file_desc == a2.m_file_desc;
        equal &= a1.m_page_size == a2.m_page_size;
        equal &= a1.m_map_size == a2.m_map_size;
        equal &= a1.m_data_ptr == a2.m_data_ptr;
        equal &= a1.m_data_offset == a2.m_data_offset;
        equal &= a1.m_fixed == a2.m_fixed;
        return equal;
    }
    
    template <class T1, class T2>
    bool operator!=(const MmapAllocator<T1>& a1,
                     const MmapAllocator<T2>& a2) throw()
    {
        return !(a1 == a2);    
    }

}

#endif