7 files changed, 185 insertions, 71 deletions

diff --git a/src/membuf.cpp b/src/membuf.cpp
index 6d850bf..e7c7ac8 100644
--- a/src/membuf.cpp
+++ b/src/membuf.cpp
@@ -41,9 +41,29 @@ size_t Bu::MemBuf::read( void *pBuf, size_t nBytes )
         
 size_t Bu::MemBuf::write( const void *pBuf, size_t nBytes )
 {
-        sBuf.append( (const char *)pBuf, nBytes );
-        nPos += nBytes;
-        return nBytes;
+        if( nPos == sBuf.getSize() )
+        {
+                // Easiest, just append the data.
+                sBuf.append( (const char *)pBuf, nBytes );
+                nPos += nBytes;
+                return nBytes;
+        }
+        else
+        {
+                // Trickier, we must do this in two parts, overwrite, then append
+                // Frist, overwrite.
+                int iOver = sBuf.getSize() - nPos;
+                if( iOver > nBytes )
+                        iOver = nBytes;
+                memcpy( sBuf.getStr()+nPos, pBuf, iOver );
+                // Then append
+                if( iOver < nBytes )
+                {
+                        sBuf.append( ((const char *)pBuf)+iOver, nBytes-iOver );
+                }
+                nPos += nBytes;
+                return nBytes;
+        }
 }
 
 long Bu::MemBuf::tell()
diff --git a/src/membuf.h b/src/membuf.h
index b728f38..15686e9 100644
--- a/src/membuf.h
+++ b/src/membuf.h
@@ -29,11 +29,6 @@ namespace Bu
                 virtual void close();
                 virtual size_t read( void *pBuf, size_t nBytes );
 
-                /**
-                 *@todo Allow writes at the current position, not just appending to
-                 * the current buffer.  This is a silly way to do it, but it covers all
-                 * of our bases for now.
-                 */
                 virtual size_t write( const void *pBuf, size_t nBytes );
                 using Stream::write;
                 virtual long tell();
diff --git a/src/nids.cpp b/src/nids.cpp
index a6596a2..06895ac 100644
--- a/src/nids.cpp
+++ b/src/nids.cpp
@@ -1,3 +1,10 @@
+/*
+ * Copyright (C) 2007-2008 Xagasoft, All rights reserved.
+ *
+ * This file is part of the libbu++ library and is released under the
+ * terms of the license contained in the file LICENSE.
+ */
+
 #include "bu/nids.h"
 #include "bu/stream.h"
 #include "bu/nidsstream.h"
@@ -16,30 +23,6 @@ Bu::Nids::Nids( Bu::Stream &sStore ) :
         iBlocks( 0 ),
         iBlockStart( -1 )
 {
-        printf("blockUnused = %u\n", blockUnused );
-        printf("Stream caps:\n"
-                "  canRead:    %s\n"
-                "  canWrite:   %s\n"
-                "  isReadable: %s\n"
-                "  isWritable: %s\n"
-                "  isSeekable: %s\n"
-                "  isBlocking: %s\n"
-                "  isEOS:      %s\n"
-                "  isOpen:     %s\n",
-                sStore.canRead()?"yes":"no",
-                sStore.canWrite()?"yes":"no",
-                sStore.isReadable()?"yes":"no",
-                sStore.isWritable()?"yes":"no",
-                sStore.isSeekable()?"yes":"no",
-                sStore.isBlocking()?"yes":"no",
-                sStore.isEOS()?"yes":"no",
-                sStore.isOpen()?"yes":"no"
-                );
-        printf("sizeof(Block) = %db\n", sizeof(Block) );
-        printf("Magic:  %02X%02X%02X%02X\n",
-                NIDS_MAGIC_CODE[0], NIDS_MAGIC_CODE[1],
-                NIDS_MAGIC_CODE[2], NIDS_MAGIC_CODE[3]
-                );
 }
 
 Bu::Nids::~Nids()
@@ -49,6 +32,7 @@ Bu::Nids::~Nids()
 void Bu::Nids::initialize()
 {
         unsigned char buf[4];
+        int iUsed;
         if( sStore.read( buf, 4 ) < 4 )
                 throw NidsException("Input stream appears to be empty.");
         if( memcmp( buf, NIDS_MAGIC_CODE, 4 ) )
@@ -56,8 +40,31 @@ void Bu::Nids::initialize()
                 throw NidsException(
                         "Stream does not appear to be a valid NIDS format.");
         }
-
-
+        sStore.read( buf, 2 );
+        if( buf[0] != 0 )
+                throw NidsException(
+                        "We can only handle version 0 for now.");
+        if( buf[1] != 4 )
+                throw NidsException(
+                        "We can only handle 4-byte words at the moment.");
+        sStore.read( &iBlockSize, 4 );
+        sStore.read( &iBlocks, 4 );
+        sStore.read( &iUsed, 4 ); // number of used blocks...
+        sStore.seek( 7 ); // skip reserved space.
+        iBlockStart = sStore.tell();
+
+        //printf("%d blocks, %db each, %db block offset\n",
+        //      iBlocks, iBlockSize, iBlockStart );
+
+        bsBlockUsed.setSize( iBlocks, true );
+        Block bTmp;
+        for( int j = 0; j < iBlocks; j++ )
+        {
+                sStore.seek( iBlockStart+iBlockSize*j );
+                sStore.read( &bTmp, sizeof(bTmp) );
+                if( bTmp.uFirstBlock != blockUnused )
+                        bsBlockUsed.setBit( j );
+        }
 }
 
 void Bu::Nids::initialize( int iBlockSize, int iPreAllocate )
@@ -88,11 +95,11 @@ void Bu::Nids::initialize( int iBlockSize, int iPreAllocate )
         this->iBlockSize = iBlockSize;
         this->iBlocks = iPreAllocate;
         this->iBlockStart = sStore.tell();
-        printf("iBlockStart = %d\n", this->iBlockStart );
+        //printf("iBlockStart = %d\n", this->iBlockStart );
         bsBlockUsed.setSize( iPreAllocate, true );
 
-        printf("%d blocks, %db each, %db block offset\n",
-                iBlocks, iBlockSize, iBlockStart );
+        //printf("%d blocks, %db each, %db block offset\n",
+        //      iBlocks, iBlockSize, iBlockStart );
 
         Block *block = (Block *)new char[iBlockSize];
         memset( block, 0, iBlockSize );
@@ -104,6 +111,26 @@ void Bu::Nids::initialize( int iBlockSize, int iPreAllocate )
         delete[] (char *)block;
 }
 
+void Bu::Nids::initBlock( uint32_t uPos, uint32_t uFirstBlock,
+        uint32_t uPrevBlock, bool bNew )
+{
+        Block b = { uPos, blockUnused, uPrevBlock, 0, 0, { } };
+        if( uFirstBlock != blockUnused )
+                b.uFirstBlock = uFirstBlock;
+        bsBlockUsed.setBit( uPos );
+        sStore.setPos( iBlockStart+(iBlockSize*uPos) );
+        sStore.write( &b, sizeof(Block) );
+        if( bNew )
+        {
+                // It's a new one, at the end, write some zeros.
+                int iSize = iBlockSize-sizeof(Block);
+                char *buf = new char[iSize];
+                memset( buf, 0, iSize );
+                sStore.write( buf, iSize );
+                delete[] buf;
+        }
+}
+
 uint32_t Bu::Nids::createBlock( uint32_t uFirstBlock, uint32_t uPrevBlock,
         int /*iPreAllocate*/ )
 {
@@ -111,16 +138,15 @@ uint32_t Bu::Nids::createBlock( uint32_t uFirstBlock, uint32_t uPrevBlock,
         {
                 if( !bsBlockUsed.getBit( j ) )
                 {
-                        Block b = { j, blockUnused, uPrevBlock, 0, 0, { } };
-                        if( uFirstBlock != blockUnused )
-                                b.uFirstBlock = uFirstBlock;
-                        bsBlockUsed.setBit( j );
-                        sStore.setPos( iBlockStart+(iBlockSize*j) );
-                        sStore.write( &b, sizeof(b) );
+                        initBlock( j, uFirstBlock, uPrevBlock );
                         return j;
                 }
         }
-        return blockUnused;
+        // Oh, we don't have any blocks left...allocate a new one.
+        iBlocks++;
+        bsBlockUsed.setSize( iBlocks, false );
+        initBlock( iBlocks-1, uFirstBlock, uPrevBlock, true );
+        return iBlocks-1;
 }
 
 int Bu::Nids::createStream( int iPreAllocate )
@@ -169,15 +195,22 @@ void Bu::Nids::updateStreamSize( uint32_t uIndex, uint32_t uSize )
 }
 
 uint32_t Bu::Nids::getNextBlock( uint32_t uIndex,
-        struct Bu::Nids::Block *pBlock )
+        struct Bu::Nids::Block *pBlock, bool bCreate )
 {
         uint32_t uNew;
         if( pBlock->uNextBlock == blockUnused )
         {
-                uNew = createBlock( pBlock->uFirstBlock, uIndex );
-                sStore.setPos( iBlockStart + (iBlockSize*uIndex)+1*4 );
-                sStore.write( &uNew, 4 );
-                getBlock( uNew, pBlock );
+                if( bCreate )
+                {
+                        uNew = createBlock( pBlock->uFirstBlock, uIndex );
+                        sStore.setPos( iBlockStart + (iBlockSize*uIndex)+1*4 );
+                        sStore.write( &uNew, 4 );
+                        getBlock( uNew, pBlock );
+                }
+                else
+                {
+                        throw Bu::NidsException("Reached end of stream.");
+                }
         }
         else
         {
diff --git a/src/nids.h b/src/nids.h
index cdefdcb..d297fa1 100644
--- a/src/nids.h
+++ b/src/nids.h
@@ -1,3 +1,10 @@
+/*
+ * Copyright (C) 2007-2008 Xagasoft, All rights reserved.
+ *
+ * This file is part of the libbu++ library and is released under the
+ * terms of the license contained in the file LICENSE.
+ */
+
 #ifndef BU_NIDS_H
 #define BU_NIDS_H
 
@@ -75,12 +82,17 @@ namespace Bu
                         blockUnused     =       0xFFFFFFFFUL
                 };
 
+                void initBlock( uint32_t uPos, uint32_t uFirstBlock,
+                        uint32_t uPrevBlock, bool bNew=false );
                 uint32_t createBlock( uint32_t uFirstBlock, uint32_t uPrevBlock,
                         int iPreAllocate=1 );
                 void getBlock( uint32_t uIndex, struct Nids::Block *pBlock );
                 void setBlock( uint32_t uIndex, struct Nids::Block *pBlock );
                 void updateStreamSize( uint32_t uIndex, uint32_t uSize );
-                uint32_t getNextBlock( uint32_t uIndex, struct Nids::Block *pBlock );
+                uint32_t getNextBlock( uint32_t uIndex, struct Nids::Block *pBlock,
+                        bool bCreate=true);
+
+                // Block allocation routines
                 Block *newBlock();
                 void deleteBlock( Block *pBlock );
 
diff --git a/src/nidsstream.cpp b/src/nidsstream.cpp
index 740dc1f..4f692c6 100644
--- a/src/nidsstream.cpp
+++ b/src/nidsstream.cpp
@@ -9,7 +9,7 @@ Bu::NidsStream::NidsStream( Nids &rNids, uint32_t uStream ) :
         uBlockSize( rNids.iBlockSize-sizeof(Nids::Block) ),
         uPos( 0 )
 {
-        printf("NidsStream::allocated\n");
+        //printf("NidsStream::allocated\n");
         pCurBlock = rNids.newBlock();
         rNids.getBlock( uStream, pCurBlock );
         uSize = pCurBlock->uBytesUsed;
@@ -34,6 +34,7 @@ Bu::NidsStream::NidsStream( const Bu::NidsStream &rSrc ) :
 
 Bu::NidsStream::~NidsStream()
 {
+        //printf("Destroying stream?\n");
         rNids.updateStreamSize( uStream, uSize );
         rNids.deleteBlock( pCurBlock );
 }
@@ -47,20 +48,35 @@ size_t Bu::NidsStream::read( void *pBuf, size_t nBytes )
         if( uPos%uBlockSize+nBytes < uBlockSize )
         {
                 size_t iRead = nBytes;
-                if( iRead > pCurBlock->uBytesUsed )
-                        iRead = pCurBlock->uBytesUsed;
+                if( iRead > pCurBlock->uBytesUsed-(uPos%uBlockSize) )
+                        iRead = pCurBlock->uBytesUsed-(uPos%uBlockSize);
                 memcpy( pBuf, pCurBlock->pData+(uPos%uBlockSize), iRead );
-                uPos += nBytes;
-                printf("a: block %u = %ub (%ub total)\n",
-                        uCurBlock, pCurBlock->uBytesUsed, uSize );
+                uPos += iRead;
+                //printf("a: block %u = %ub (%ub total)\n",
+                //      uCurBlock, pCurBlock->uBytesUsed, uSize );
                 return iRead;
         }
         else
         {
-                //size_t iTotal = 0;
+                size_t nTotal = 0;
                 for(;;)
                 {
-                        
+                        uint32_t iRead = nBytes;
+                        if( iRead > uBlockSize-(uPos%uBlockSize) )
+                                iRead = uBlockSize-(uPos%uBlockSize);
+                        if( iRead > pCurBlock->uBytesUsed-(uPos%uBlockSize) )
+                                iRead = pCurBlock->uBytesUsed-(uPos%uBlockSize);
+                        memcpy( ((char *)pBuf)+nTotal,
+                                pCurBlock->pData+(uPos%uBlockSize), iRead );
+                        uPos += iRead;
+                        nBytes -= iRead;
+                        nTotal += iRead;
+                        //printf("r: block %u = %ub/%ub (%ub total)\n",
+                        //      uCurBlock, iRead, pCurBlock->uBytesUsed, uSize );
+                        if( nBytes == 0 || uPos == uSize )
+                                return nTotal;
+                        if( nTotal%uBlockSize == 0 )
+                                uCurBlock = rNids.getNextBlock( uCurBlock, pCurBlock, false );
                 }
         }
         return 0;
@@ -75,8 +91,8 @@ size_t Bu::NidsStream::write( const void *pBuf, size_t nBytes )
                 rNids.setBlock( uCurBlock, pCurBlock );
                 uPos += nBytes;
                 uSize += nBytes;
-                printf("a: block %u = %ub (%ub total)\n",
-                        uCurBlock, pCurBlock->uBytesUsed, uSize );
+                //printf("a: block %u = %ub (%ub total)\n",
+                //      uCurBlock, pCurBlock->uBytesUsed, uSize );
                 return nBytes;
         }
         else
@@ -95,8 +111,8 @@ size_t Bu::NidsStream::write( const void *pBuf, size_t nBytes )
                         uPos += uNow;
                         nTotal += uNow;
                         nBytes -= uNow;
-                        printf("b: block %u = %ub (%ub total)\n",
-                                uCurBlock, pCurBlock->uBytesUsed, uSize );
+                        //printf("b: block %u = %ub (%ub total)\n",
+                        //      uCurBlock, pCurBlock->uBytesUsed, uSize );
                         if( nBytes == 0 )
                                 return nTotal;
                         if( pCurBlock->uBytesUsed == uBlockSize )
diff --git a/src/tests/nids.cpp b/src/tests/nids.cpp
index f50fde5..d531280 100644
--- a/src/tests/nids.cpp
+++ b/src/tests/nids.cpp
@@ -4,22 +4,44 @@
 
 int main( int argc, char *argv[] )
 {
-        if( argc < 2 )
+        if( argc < 3 )
         {
-                printf("usage: %s <output>\n\n", argv[0] );
+                printf("usage: %s [r|w] <output>\n\n", argv[0] );
                 return 1;
         }
 
-        Bu::File fOut( argv[1], Bu::File::ReadWrite );
-        Bu::Nids n( fOut );
+        if( argv[1][0] == 'w' )
+        {
+                Bu::File fOut( argv[2],
+                        Bu::File::ReadWrite|Bu::File::Create|Bu::File::Truncate );
+                Bu::Nids n( fOut );
 
-//      n.initialize( 120, 5 );
+                n.initialize( 120, 1 );
+                Bu::NidsStream s = n.openStream( n.createStream() );
 
-        Bu::NidsStream s = n.openStream( n.createStream() );
+                Bu::FString sBuf( 350 );
+                memset( sBuf.getStr(), 'a', 350 );
+                s.write( sBuf );
+        }
+        else if( argv[1][0] == 'r' )
+        {
+                Bu::File fOut( argv[2], Bu::File::Read );
+                Bu::Nids n( fOut );
 
-//      Bu::FString sBuf( 350 );
-//      memset( sBuf.getStr(), 'a', 350 );
-//      s.write( sBuf );
+                Bu::NidsStream s = n.openStream( 0 );
+                char buf[75];
+                for( int j = 0; j < 3; j++ )
+                {
+                        int iRead = s.read( buf, 75 );
+                        fwrite( buf, 1, iRead, stdout );
+                        printf("\n(read %d chars)\n", iRead );
+                }
+                printf("Position:  %ld\n", s.tell() );
+        }
+        else
+        {
+                printf("r or w, those are your choices.\n");
+        }
 
         return 0;
 }
diff --git a/src/unit/membuf.cpp b/src/unit/membuf.cpp
index 3aebc4d..dc02aa3 100644
--- a/src/unit/membuf.cpp
+++ b/src/unit/membuf.cpp
@@ -15,6 +15,7 @@ public:
         {
                 setName("MemBuf");
                 addTest( Unit::testWriteRead01 );
+                addTest( Unit::testOverwrite1 );
         }
 
         virtual ~Unit()
@@ -39,6 +40,21 @@ public:
                 unitTest( mb.read( buf, 7 ) == 3 );
                 unitTest( !strncmp( buf, "eFG", 3 ) );
         }
+
+        void testOverwrite1()
+        {
+                Bu::MemBuf mb;
+                unitTest( mb.write("0123456789") == 10 );
+                mb.setPos( 4 );
+                unitTest( mb.write("-5-") == 3 );
+                mb.setPos( 9 );
+                mb.write("Hey!!!");
+                unitTest( mb.tell() == 15 );
+                char buf[50];
+                mb.setPos( 0 );
+                buf[mb.read( buf, 50 )] = '\0';
+                unitTest( !strcmp( buf, "0123-5-78Hey!!!" ) );
+        }
 };
 
 int main( int argc, char *argv[] ){ return Unit().run( argc, argv ); }