import * as assert from 'assert' ; import { IEndpoint } from '../types' ; import { Buffer } from 'buffer' ; import { isUndefined } from 'lodash' ; import { copy , resize } from './array' ; type RemovedCallback = ( endpoint : IEndpoint , buffer : Buffer , bufferSize : number ) => void ; type MatchCallback = ( endpoint : IEndpoint ) => void ; export default class MultiTrie { endpoints ? : Set < IEndpoint > ; minCharacter : number ; count : number ; liveNodes : number ; next : Array < MultiTrie | null > ; constructor ( ) { this . minCharacter = 0 ; this . count = 0 ; this . liveNodes = 0 ; this . next = [ ] ; } get isRedundant ( ) : boolean { return isUndefined ( this . endpoints ) && this . liveNodes === 0 ; } add ( prefix , start , size , endpoint ) { return this . addHelper ( prefix , start , size , endpoint ) ; } private addHelper ( prefix , start , size , endpoint ) { if ( size === 0 ) { let result = isUndefined ( this . endpoints ) ; if ( isUndefined ( this . endpoints ) ) this . endpoints = new Set < IEndpoint > ( ) ; this . endpoints . add ( endpoint ) ; return result ; } const currentCharacter = prefix . readUInt8 ( start ) ; if ( currentCharacter < this . minCharacter || currentCharacter >= this . minCharacter + this . count ) { if ( this . count === 0 ) { this . minCharacter = currentCharacter ; this . count = 1 ; this . next = Array ( 1 ) . fill ( null ) ; } else if ( this . count === 1 ) { let oldc = this . minCharacter ; const oldp = this . next [ 0 ] ; this . count = ( this . minCharacter < currentCharacter ? currentCharacter - this . minCharacter : this . minCharacter - currentCharacter ) + 1 ; this . next = Array ( this . count ) . fill ( null ) ; this . minCharacter = Math . min ( this . minCharacter , currentCharacter ) ; this . next [ oldc - this . minCharacter ] = oldp ; } else if ( this . minCharacter < currentCharacter ) { this . count = currentCharacter - this . minCharacter + 1 ; this . next = resize ( this . next , this . count , true ) ; } else { this . count = ( this . minCharacter + this . count ) - currentCharacter ; this . next = resize ( this . next , this . count , false ) ; this . minCharacter = currentCharacter ; } } if ( this . next [ currentCharacter - this . minCharacter ] === null ) { this . next [ currentCharacter - this . minCharacter ] = new MultiTrie ( ) ; this . liveNodes ++ ; } return this . next [ currentCharacter - this . minCharacter ] . addHelper ( prefix , start + 1 , size - 1 , endpoint ) ; } public removeEndpoint ( endpoint , func ) { return this . removeEndpointHelper ( endpoint , Buffer . alloc ( 0 ) , 0 , 0 , func ) ; } private removeEndpointHelper ( endpoint , buffer , bufferSize , maxBufferSize , func ) { if ( this . endpoints && this . endpoints . delete ( endpoint ) && this . endpoints . size === 0 ) { func ( endpoint , buffer , bufferSize ) ; this . endpoints = undefined ; } if ( bufferSize >= maxBufferSize ) { maxBufferSize = bufferSize + 256 ; const newBuffer = Buffer . alloc ( maxBufferSize , 0 ) ; buffer . copy ( newBuffer ) ; buffer = newBuffer ; } if ( this . count === 0 ) return true ; if ( this . count === 1 ) { buffer . writeUInt8 ( this . minCharacter , bufferSize ) ; bufferSize ++ ; this . next [ 0 ] . removeEndpointHelper ( endpoint , buffer , bufferSize , maxBufferSize , func ) ; if ( this . next [ 0 ] . isRedundant ) { this . next = [ ] ; this . count = 0 ; this . liveNodes -- ; } return true ; } let newMin = this . minCharacter + this . count - 1 ; let newMax = this . minCharacter ; for ( let currentCharacter = 0 ; currentCharacter != this . count ; currentCharacter ++ ) { buffer . writeUInt8 ( this . minCharacter + currentCharacter , bufferSize ) ; const next = this . next [ currentCharacter ] ; if ( next ) { next . removeEndpointHelper ( endpoint , buffer , bufferSize + 1 , maxBufferSize , func ) ; if ( next . isRedundant ) { this . next [ currentCharacter ] = null ; this . liveNodes -- ; } else { if ( currentCharacter + this . minCharacter < newMin ) newMin = currentCharacter + this . minCharacter ; if ( currentCharacter + this . minCharacter > newMax ) newMax = currentCharacter + this . minCharacter ; } } } if ( this . liveNodes === 0 ) { this . next = [ ] ; this . count = 0 ; } else if ( this . liveNodes === 1 ) { const node = this . next [ newMin - this . minCharacter ] ; assert ( node ) ; this . next = [ node ] ; this . count = 1 ; this . minCharacter = newMin ; } else if ( this . liveNodes > 1 && ( newMin > this . minCharacter || newMax < this . minCharacter + this . count - 1 ) ) { let oldTable = this . next ; this . count = newMax - newMin + 1 ; this . next = Array ( this . count ) . fill ( null ) ; copy ( oldTable , ( newMin - this . minCharacter ) , this . next , 0 , this . count ) ; this . minCharacter = newMin ; } return true ; } public remove ( prefix , start , size , endpoint ) { if ( size === 0 ) { if ( this . endpoints ) { const erased = this . endpoints . delete ( endpoint ) ; assert ( erased ) ; if ( this . endpoints . size === 0 ) this . endpoints = undefined ; } return ! this . endpoints ; } const currentCharacter = prefix . readUInt8 ( start ) ; if ( this . count == 0 || currentCharacter < this . minCharacter || currentCharacter >= this . minCharacter + this . count ) return false ; const nextNode = this . count == 1 ? this . next [ 0 ] : this . next [ currentCharacter - this . minCharacter ] ; if ( nextNode === null ) return false ; let ret = nextNode . remove ( prefix , start + 1 , size - 1 , endpoint ) ; if ( nextNode . isRedundant ) { assert ( this . count > 0 ) ; if ( this . count === 1 ) { this . next = [ ] ; this . count = 0 ; this . liveNodes -- ; } else { this . next [ currentCharacter - this . minCharacter ] = null ; this . liveNodes -- ; if ( this . liveNodes === 1 ) { let i = 0 ; for ( ; i < this . count ; i ++ ) { if ( this . next [ i ] !== null ) break ; } this . minCharacter += i ; this . count = 1 ; const old = this . next [ i ] ; this . next = [ old ] ; } else if ( currentCharacter == this . minCharacter ) { let i = 1 ; for ( ; i < this . count ; i ++ ) { if ( this . next [ i ] !== null ) break ; } this . minCharacter += i ; this . count -= i ; this . next = resize ( this . next , this . count , false ) ; } else if ( currentCharacter === this . minCharacter + this . count - 1 ) { let i = 1 ; for ( ; i < this . count ; i ++ ) { if ( this . next [ this . count - 1 - i ] !== null ) break ; } this . count -= i ; this . next = resize ( this . next , this . count , true ) ; } } } return ret ; } match ( data , offset , size , func ) { let current = this ; let index = offset ; while ( true ) { if ( current . endpoints ) current . endpoints . forEach ( e => func ( e ) ) ; if ( size === 0 ) break ; if ( current . count === 0 ) break ; const c = data . readUInt8 ( index ) ; if ( current . count === 1 ) { if ( c != current . minCharacter ) break ; current = current . next [ 0 ] ; index ++ ; size -- ; continue ; } if ( c < current . minCharacter || c >= current . minCharacter + current . count ) break ; if ( current . next [ c - current . minCharacter ] === null ) break ; current = current . next [ c - current . minCharacter ] ; index ++ ; size -- ; } } }