extra/bson/writer/writer.factor

   1 ! Copyright (C) 2008 Sascha Matzke.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: accessors assocs bson.constants byte-arrays byte-vectors
   4 calendar fry io io.binary io.encodings io.encodings.binary
   5 io.encodings.utf8 io.streams.byte-array kernel math math.parser
   6 namespaces quotations sequences sequences.private serialize strings
   7 words combinators.short-circuit literals ;
   8
   9
  10 IN: bson.writer
  11
  12 #! Writes the object out to a byte-vector in BSON format
  13
  14 <PRIVATE
  15
  16 SYMBOL: shared-buffer
  17
  18 CONSTANT: INT32-SIZE 4
  19 CONSTANT: CHAR-SIZE 1
  20 CONSTANT: INT64-SIZE 8
  21
  22 : (buffer) ( -- buffer )
  23     shared-buffer get
  24     [ 8192 <byte-vector> [ shared-buffer set ] keep ] unless* ; inline
  25
  26 : >le-stream ( x n -- )
  27     ! >le write
  28     swap '[ _ swap nth-byte 0 B{ 0 }
  29             [ set-nth-unsafe ] keep write ] each
  30             ; inline
  31
  32 PRIVATE>
  33
  34 : reset-buffer ( buffer -- )
  35     0 >>length drop ; inline
  36
  37 : ensure-buffer ( -- )
  38     (buffer) drop ; inline
  39
  40 : with-buffer ( quot -- byte-vector )
  41     [ (buffer) [ reset-buffer ] keep dup ] dip
  42     with-output-stream* dup encoder? [ stream>> ] when ; inline
  43
  44 : with-length ( quot: ( -- ) -- bytes-written start-index )
  45     [ (buffer) [ length ] keep ] dip call
  46     length swap [ - ] keep ; inline
  47
  48 : with-length-prefix ( quot: ( -- ) -- )
  49     [ B{ 0 0 0 0 } write ] prepose with-length
  50     [ INT32-SIZE >le ] dip (buffer)
  51     '[ _ over [ nth-unsafe ] dip _ + _ set-nth-unsafe ]
  52     [ INT32-SIZE ] dip each-integer ; inline
  53
  54 : with-length-prefix-excl ( quot: ( -- ) -- )
  55     [ B{ 0 0 0 0 } write ] prepose with-length
  56     [ INT32-SIZE - INT32-SIZE >le ] dip (buffer)
  57     '[ _ over [ nth-unsafe ] dip _ + _ set-nth-unsafe ]
  58     [ INT32-SIZE ] dip each-integer ; inline
  59
  60 <PRIVATE
  61
  62 GENERIC: bson-type? ( obj -- type ) foldable flushable
  63 GENERIC: bson-write ( obj -- )
  64
  65 M: t bson-type? ( boolean -- type ) drop T_Boolean ;
  66 M: f bson-type? ( boolean -- type ) drop T_Boolean ;
  67
  68 M: real bson-type? ( real -- type ) drop T_Double ;
  69 M: word bson-type? ( word -- type ) drop T_String ;
  70 M: tuple bson-type? ( tuple -- type ) drop T_Object ;
  71 M: sequence bson-type? ( seq -- type ) drop T_Array ;
  72 M: string bson-type? ( string -- type ) drop T_String ;
  73 M: integer bson-type? ( integer -- type ) drop T_Integer ;
  74 M: assoc bson-type? ( assoc -- type ) drop T_Object ;
  75 M: timestamp bson-type? ( timestamp -- type ) drop T_Date ;
  76 M: mdbregexp bson-type? ( regexp -- type ) drop T_Regexp ;
  77
  78 M: oid bson-type? ( word -- type ) drop T_OID ;
  79 M: objref bson-type? ( objref -- type ) drop T_Binary ;
  80 M: quotation bson-type? ( quotation -- type ) drop T_Binary ;
  81 M: byte-array bson-type? ( byte-array -- type ) drop T_Binary ;
  82
  83 : write-utf8-string ( string -- )
  84     output-stream get utf8 <encoder> stream-write ; inline
  85
  86 : write-byte ( byte -- ) CHAR-SIZE >le-stream ; inline
  87 : write-int32 ( int -- ) INT32-SIZE >le-stream ; inline
  88 : write-double ( real -- ) double>bits INT64-SIZE >le-stream ; inline
  89 : write-cstring ( string -- ) write-utf8-string 0 write-byte ; inline
  90 : write-longlong ( object -- ) INT64-SIZE >le-stream ; inline
  91
  92 : write-eoo ( -- ) T_EOO write-byte ; inline
  93 : write-type ( obj -- obj ) [ bson-type? write-byte ] keep ; inline
  94 : write-pair ( name object -- ) write-type [ write-cstring ] dip bson-write ; inline
  95
  96 M: f bson-write ( f -- )
  97     drop 0 write-byte ;
  98
  99 M: t bson-write ( t -- )
 100     drop 1 write-byte ;
 101
 102 M: string bson-write ( obj -- )
 103     '[ _ write-cstring ] with-length-prefix-excl ;
 104
 105 M: integer bson-write ( num -- )
 106     write-int32 ;
 107
 108 M: real bson-write ( num -- )
 109     >float write-double ;
 110
 111 M: timestamp bson-write ( timestamp -- )
 112     timestamp>millis write-longlong ;
 113
 114 M: byte-array bson-write ( binary -- )
 115     [ length write-int32 ] keep
 116     T_Binary_Bytes write-byte
 117     write ;
 118
 119 M: quotation bson-write ( quotation -- )
 120     object>bytes [ length write-int32 ] keep
 121     T_Binary_Function write-byte
 122     write ;
 123
 124 M: oid bson-write ( oid -- )
 125     [ a>> write-longlong ] [ b>> write-int32 ] bi ;
 126
 127 M: objref bson-write ( objref -- )
 128     [ binary ] dip
 129     '[ _
 130        [ ns>> write-cstring ]
 131        [ objid>> write-cstring ] bi ] with-byte-writer
 132     [ length write-int32 ] keep
 133     T_Binary_Custom write-byte write ;
 134
 135 M: mdbregexp bson-write ( regexp -- )
 136    [ regexp>> write-cstring ]
 137    [ options>> write-cstring ] bi ;
 138
 139 M: sequence bson-write ( array -- )
 140     '[ _ [ [ write-type ] dip number>string
 141            write-cstring bson-write ] each-index
 142        write-eoo ] with-length-prefix ;
 143
 144 : write-oid ( assoc -- )
 145     [ MDB_OID_FIELD ] dip at
 146     [ [ MDB_OID_FIELD ] dip write-pair ] when* ; inline
 147
 148 : skip-field? ( name -- boolean )
 149    { $[ MDB_OID_FIELD MDB_META_FIELD ] } member? ; inline
 150
 151 M: assoc bson-write ( assoc -- )
 152     '[ _  [ write-oid ] keep
 153        [ over skip-field? [ 2drop ] [ write-pair ] if ] assoc-each
 154        write-eoo ] with-length-prefix ;
 155
 156 M: word bson-write name>> bson-write ;
 157
 158 PRIVATE>
 159
 160 : assoc>bv ( assoc -- byte-vector )
 161     [ '[ _ bson-write ] with-buffer ] with-scope ; inline
 162
 163 : assoc>stream ( assoc -- )
 164     bson-write ; inline
 165
 166 : mdb-special-value? ( value -- ? )
 167    { [ timestamp? ] [ quotation? ] [ mdbregexp? ]
 168      [ oid? ] [ byte-array? ] } 1|| ;