Fix atomic fetch operations

chibicc.h

@@ -282,6 +282,9 @@ struct Node {
   Obj *atomic_addr;
   Node *atomic_expr;
 
+  // Atomic fetch operation
+  bool atomic_fetch;
+
   // Variable
   Obj *var;
 

include/stdatomic.h

@@ -34,17 +34,17 @@ typedef enum {
 #define atomic_load_explicit(addr, order) (*(addr))
 #define atomic_store_explicit(addr, val, order) (*(addr) = (val))
 
-#define atomic_fetch_add(obj, val) (*(obj) += (val))
-#define atomic_fetch_sub(obj, val) (*(obj) -= (val))
-#define atomic_fetch_or(obj, val) (*(obj) |= (val))
-#define atomic_fetch_xor(obj, val) (*(obj) ^= (val))
-#define atomic_fetch_and(obj, val) (*(obj) &= (val))
+#define atomic_fetch_add(obj, val) __builtin_atomic_fetch_op(obj, val, 0)
+#define atomic_fetch_sub(obj, val) __builtin_atomic_fetch_op(obj, val, 1)
+#define atomic_fetch_or(obj, val) __builtin_atomic_fetch_op(obj, val, 2)
+#define atomic_fetch_xor(obj, val) __builtin_atomic_fetch_op(obj, val, 3)
+#define atomic_fetch_and(obj, val) __builtin_atomic_fetch_op(obj, val, 4)
 
-#define atomic_fetch_add_explicit(obj, val, order) (*(obj) += (val))
-#define atomic_fetch_sub_explicit(obj, val, order) (*(obj) -= (val))
-#define atomic_fetch_or_explicit(obj, val, order) (*(obj) |= (val))
-#define atomic_fetch_xor_explicit(obj, val, order) (*(obj) ^= (val))
-#define atomic_fetch_and_explicit(obj, val, order) (*(obj) &= (val))
+#define atomic_fetch_add_explicit(obj, val, order) atomic_fetch_add(obj, val)
+#define atomic_fetch_sub_explicit(obj, val, order) atomic_fetch_sub(obj, val)
+#define atomic_fetch_or_explicit(obj, val, order) atomic_fetch_or(obj, val)
+#define atomic_fetch_xor_explicit(obj, val, order) atomic_fetch_xor(obj, val)
+#define atomic_fetch_and_explicit(obj, val, order) atomic_fetch_and(obj, val)
 
 #define atomic_compare_exchange_weak(p, old, new) \
   __builtin_compare_and_swap((p), (old), (new))

parse.c

@@ -2068,14 +2068,15 @@ static Node *to_assign(Node *binary) {
   //   } while (!atomic_compare_exchange_strong(addr, &old, new));
   //   new;
   // })
-  if (binary->lhs->ty->is_atomic) {
+  if (binary->lhs->ty->is_atomic || binary->atomic_fetch) {
     Node head = {};
     Node *cur = &head;
 
     Obj *addr = new_lvar("", pointer_to(binary->lhs->ty));
     Obj *val = new_lvar("", binary->rhs->ty);
     Obj *old = new_lvar("", binary->lhs->ty);
     Obj *new = new_lvar("", binary->lhs->ty);
+    Obj *ret = binary->atomic_fetch ? old : new;
 
     cur = cur->next =
       new_unary(ND_EXPR_STMT,
@@ -2114,7 +2115,7 @@ static Node *to_assign(Node *binary) {
     loop->cond = new_unary(ND_NOT, cas, tok);
 
     cur = cur->next = loop;
-    cur = cur->next = new_unary(ND_EXPR_STMT, new_var_node(new, tok), tok);
+    cur = cur->next = new_unary(ND_EXPR_STMT, new_var_node(ret, tok), tok);
 
     Node *node = new_node(ND_STMT_EXPR, tok);
     node->body = head.next;
@@ -3079,6 +3080,32 @@ static Node *primary(Token **rest, Token *tok) {
     return node;
   }
 
+  if (equal(tok, "__builtin_atomic_fetch_op")) {
+    tok = skip(tok->next, "(");
+    Node *obj = new_unary(ND_DEREF, assign(&tok, tok), tok);
+    tok = skip(tok, ",");
+    Node *val = assign(&tok, tok);
+    tok = skip(tok, ",");
+    Node *node;
+
+    if (equal(tok, "0"))
+      node = new_add(obj, val, tok);
+    else if (equal(tok, "1"))
+      node = new_sub(obj, val, tok);
+    else if (equal(tok, "2"))
+      node = new_binary(ND_BITOR, obj, val, tok);
+    else if (equal(tok, "3"))
+      node = new_binary(ND_BITXOR, obj, val, tok);
+    else if (equal(tok, "4"))
+      node = new_binary(ND_BITAND, obj, val, tok);
+    else
+      error_tok(tok, "invalid fetch operator");
+
+    node->atomic_fetch = true;
+    *rest = skip(tok->next, ")");
+    return to_assign(node);
+  }
+
   if (tok->kind == TK_IDENT) {
     // Variable or enum constant
     VarScope *sc = find_var(tok);

test/atomic.c

@@ -52,12 +52,30 @@ static int add_millions(void) {
   return x;
 }
 
+static void fetch_ops(void) {
+  _Atomic int x = 0;
+
+  ASSERT(0, atomic_fetch_add(&x, 17));
+  ASSERT(17, atomic_fetch_add(&x, 10));
+  ASSERT(27, atomic_fetch_add(&x, 3));
+  ASSERT(30, atomic_fetch_sub(&x, 17));
+  ASSERT(13, atomic_fetch_sub(&x, 13));
+
+  ASSERT(0, atomic_fetch_or(&x, 0xf0));
+  ASSERT(0xf0, atomic_fetch_or(&x, 0x0f));
+  ASSERT(0xff, atomic_fetch_and(&x, 0x0f));
+  ASSERT(0x0f, atomic_fetch_xor(&x, 0xff));
+  ASSERT(0xf0, atomic_fetch_add(&x, 0));
+}
+
 int main() {
   ASSERT(6*1000*1000, add_millions());
 
   ASSERT(3, ({ int x=3; atomic_exchange(&x, 5); }));
   ASSERT(5, ({ int x=3; atomic_exchange(&x, 5); x; }));
 
+  fetch_ops();
+
   printf("OK\n");
   return 0;
 }