io.sockets: receive-unsafe and receive-into

Analogous to read-unsafe/read-into for streams (and thereby TCP sockets), provide receive-unsafe and receive-into words for datagram sockets that receive into a caller-supplied buffer.

basis/io/sockets/sockets.factor

@@ -4,11 +4,11 @@
 USING: accessors alien.c-types alien.data alien.strings arrays
 assocs byte-arrays classes classes.struct combinators
 combinators.short-circuit continuations destructors fry generic
-grouping init io.backend io.pathnames io.binary io.encodings
+grouping init io.backend io.binary io.encodings
-io.encodings.ascii io.encodings.binary io.ports
+io.encodings.ascii io.encodings.binary io.pathnames io.ports
-io.streams.duplex kernel math math.parser memoize namespaces
+io.streams.duplex kernel libc locals math math.parser memoize
-parser present sequences splitting strings summary system
+namespaces parser present sequences splitting strings summary
-vocabs.loader vocabs.parser ;
+system unix.ffi values vocabs.loader vocabs.parser ;
 IN: io.sockets
 
 << {
@@ -286,7 +286,7 @@ TUPLE: raw-port < port addr ;
 
 HOOK: (raw) io-backend ( addr -- raw )
 
-HOOK: (receive) io-backend ( datagram -- packet addrspec )
+HOOK: (receive) io-backend ( n buf datagram -- size addrspec )
 
 ERROR: invalid-port object ;
 
@@ -368,17 +368,37 @@ SYMBOL: remote-address
         >>addr
     ] with-destructors ;
 
-: receive ( datagram -- packet addrspec )
+: receive-unsafe ( n buf datagram -- count addrspec )
     check-receive
-    [ (receive) ] [ addr>> ] bi parse-sockaddr ;
+    [ (receive) ] [ addr>> ] bi parse-sockaddr ; inline
+
+CONSTANT: datagram-size 65536
+STRUCT: datagram-buf { buf uchar[datagram-size] } ;
+
+:: receive ( datagram -- packet addrspec )
+    { datagram-buf } [| buf |
+        datagram-size buf datagram
+        receive-unsafe :> ( count addrspec )
+        count [ f f ] [
+            buf swap memory>byte-array addrspec
+        ] if-zero
+    ] with-scoped-allocation ; inline
+
+:: receive-into ( buf datagram -- buf-slice addrspec )
+    buf length :> n
+    n buf datagram receive-unsafe :> ( count addrspec )
+    count [ f f ] [ drop
+        buf count head-slice addrspec
+    ] if-zero ; inline
 
 : send ( packet addrspec datagram -- )
-    check-send (send) ;
+    check-send (send) ; inline
 
 MEMO: ipv6-supported? ( -- ? )
     [ "::1" 0 <inet6> binary <server> dispose t ] [ drop f ] recover ;
 
-[ \ ipv6-supported? reset-memoized ] "io.sockets" add-startup-hook
+[ \ ipv6-supported? reset-memoized ]
+"io.sockets:ipv6-supported?" add-startup-hook
 
 GENERIC: resolve-host ( addrspec -- seq )
 

basis/io/sockets/unix/unix.factor

@@ -125,29 +125,23 @@ M: unix (datagram)
 M: unix (raw)
     [ SOCK_RAW server-socket-fd ] with-destructors ;
 
-SYMBOL: receive-buffer
+:: do-receive ( n buf port -- count sockaddr )
-
-CONSTANT: packet-size 65536
-
-[ packet-size malloc &free receive-buffer set-global ] "io.sockets.unix" add-startup-hook
-
-:: do-receive ( port -- packet sockaddr )
     port addr>> empty-sockaddr/size :> ( sockaddr len )
     port handle>> handle-fd ! s
-    receive-buffer get-global ! buf
+    buf ! buf
-    packet-size ! nbytes
+    n ! nbytes
     0 ! flags
     sockaddr ! from
     len int <ref> ! fromlen
-    recvfrom dup 0 >=
+    recvfrom sockaddr ; inline
-    [ receive-buffer get-global swap memory>byte-array sockaddr ]
-    [ drop f f ]
-    if ;
 
-M: unix (receive) ( datagram -- packet sockaddr )
+: (receive-loop) ( n buf datagram -- count sockaddr )
-    dup do-receive dup [ [ drop ] 2dip ] [
+    3dup do-receive over 0 > [ [ 3drop ] 2dip ] [
-        2drop [ +input+ wait-for-port ] [ (receive) ] bi
+        2drop [ +input+ wait-for-port ] [ (receive-loop) ] bi
-    ] if ;
+    ] if ; inline recursive
+
+M: unix (receive) ( n buf datagram -- count sockaddr )
+    (receive-loop) ;
 
 :: do-send ( packet sockaddr len socket datagram -- )
     socket handle-fd packet dup length 0 sockaddr len sendto
@@ -162,7 +156,7 @@ M: unix (receive) ( datagram -- packet sockaddr )
                 (io-error)
             ] if
         ] if
-    ] when ;
+    ] when ; inline recursive
 
 M: unix (send) ( packet addrspec datagram -- )
     [ make-sockaddr/size ] [ [ handle>> ] keep ] bi* do-send ;

basis/io/sockets/windows/windows.factor

@@ -211,19 +211,19 @@ TUPLE: WSARecvFrom-args port
        s lpBuffers dwBufferCount lpNumberOfBytesRecvd
        lpFlags lpFrom lpFromLen lpOverlapped lpCompletionRoutine ;
 
-: make-receive-buffer ( -- WSABUF )
+: make-receive-buffer ( n buf -- WSABUF )
     WSABUF malloc-struct &free
-        default-buffer-size get
+        n >>len
-        [ >>len ] [ malloc &free >>buf ] bi ; inline
+        buf >>buf ; inline
 
-: <WSARecvFrom-args> ( datagram -- WSARecvFrom )
+:: <WSARecvFrom-args> ( n buf datagram -- WSARecvFrom )
     WSARecvFrom-args new
-        swap >>port
+        datagram >>port
-        dup port>> handle>> handle>> >>s
+        datagram handle>> handle>> >>s
-        dup port>> addr>> sockaddr-size
+        datagram addr>> sockaddr-size
             [ malloc &free >>lpFrom ]
             [ malloc-int &free >>lpFromLen ] bi
-        make-receive-buffer >>lpBuffers
+        n buf make-receive-buffer >>lpBuffers
         1 >>dwBufferCount
         0 malloc-int &free >>lpFlags
         0 malloc-int &free >>lpNumberOfBytesRecvd
@@ -251,7 +251,7 @@ TUPLE: WSARecvFrom-args port
         tri memcpy
     ] bi ; inline
 
-M: windows (receive) ( datagram -- packet addrspec )
+M: windows (receive) ( n buf datagram -- packet addrspec )
     [
         <WSARecvFrom-args>
         [ call-WSARecvFrom ]