package org.jruby.ext.atomic;

import java.io.IOException;
import java.util.concurrent.atomic.AtomicReference;
import org.jruby.Ruby;
import org.jruby.RubyClass;
import org.jruby.RubyModule;
import org.jruby.RubyObject;
import org.jruby.anno.JRubyClass;
import org.jruby.anno.JRubyMethod;
import org.jruby.runtime.ObjectAllocator;
import org.jruby.runtime.ThreadContext;
import org.jruby.runtime.builtin.IRubyObject;
import org.jruby.runtime.load.Library;

/**
 * This library adds an atomic reference type to JRuby for use in the atomic
 * library. We do a native version to avoid the implicit value coercion that
 * normally happens through JI.
 * 
 * @author headius
 */
public class AtomicReferenceLibrary implements Library {
    public void load(Ruby runtime, boolean wrap) throws IOException {
        RubyModule atomicCls = runtime.getClass("Atomic");
        RubyClass jrubyRefClass = runtime.defineClassUnder("InternalReference", runtime.getObject(), JRUBYREFERENCE_ALLOCATOR, atomicCls);
        jrubyRefClass.setAllocator(JRUBYREFERENCE_ALLOCATOR);
        jrubyRefClass.defineAnnotatedMethods(JRubyReference.class);
    }
    
    private static final ObjectAllocator JRUBYREFERENCE_ALLOCATOR = new ObjectAllocator() {
        public IRubyObject allocate(Ruby runtime, RubyClass klazz) {
            return new JRubyReference(runtime, klazz);
        }
    };

    @JRubyClass(name="JRubyReference", parent="Object")
    public static class JRubyReference extends RubyObject {
        private final AtomicReference<IRubyObject> reference;

        public JRubyReference(Ruby runtime, RubyClass klass) {
            super(runtime, klass);
            reference = new AtomicReference<IRubyObject>(runtime.getNil());
        }

        @JRubyMethod
        public IRubyObject initialize(ThreadContext context) {
            Ruby runtime = context.getRuntime();
            reference.set(runtime.getNil());
            return runtime.getNil();
        }

        @JRubyMethod
        public IRubyObject initialize(ThreadContext context, IRubyObject value) {
            Ruby runtime = context.getRuntime();
            reference.set(value);
            return runtime.getNil();
        }

        @JRubyMethod(name = {"get", "value"})
        public IRubyObject get() {
            return reference.get();
        }

        @JRubyMethod(name = {"set", "value="})
        public IRubyObject set(IRubyObject newValue) {
            reference.set(newValue);
            return newValue;
        }

        @JRubyMethod
        public IRubyObject compare_and_set(ThreadContext context, IRubyObject oldValue, IRubyObject newValue) {
            return context.getRuntime().newBoolean(reference.compareAndSet(oldValue, newValue));
        }

        @JRubyMethod
        public IRubyObject get_and_set(ThreadContext context, IRubyObject newValue) {
            return reference.getAndSet(newValue);
        }
    }
}

@@ -102,18 +102,30 @@
 eference
+FieldUpdater
 ;%0A
-
 import o
@@ -1495,16 +1495,64 @@
 bject %7B%0A
+        private volatile IRubyObject reference;%0A
         
@@ -1565,16 +1565,23 @@
 e final 
+static 
 AtomicRe
@@ -1591,22 +1591,153 @@
 ence
-%3CIRubyObject%3E 
+FieldUpdater%3CJRubyReference, IRubyObject%3E UPDATER =%0A            AtomicReferenceFieldUpdater.newUpdater(JRubyReference.class, IRubyObject.class, %22
 refe
@@ -1741,16 +1741,18 @@
 eference
+%22)
 ;%0A%0A     
@@ -1872,41 +1872,8 @@
 e = 
-new AtomicReference%3CIRubyObject%3E(
 runt
@@ -1876,33 +1876,32 @@
 runtime.getNil()
-)
 ;%0A        %7D%0A%0A   
@@ -2042,30 +2042,34 @@
         
-reference.set(
+UPDATER.set(this, 
 runtime.
@@ -2291,30 +2291,34 @@
         
-reference.set(
+UPDATER.set(this, 
 value);%0A
@@ -2469,22 +2469,24 @@
 urn 
-reference
+UPDATER
 .get(
+this
 );%0A 
@@ -2613,22 +2613,26 @@
     
-reference.set(
+UPDATER.set(this, 
 newV
@@ -2862,25 +2862,23 @@
 Boolean(
-reference
+UPDATER
 .compare
@@ -2884,16 +2884,22 @@
 eAndSet(
+this, 
 oldValue
@@ -3049,25 +3049,23 @@
  return 
-reference
+UPDATER
 .getAndS
@@ -3067,16 +3067,22 @@
 tAndSet(
+this, 
 newValue

Use AtomicReferenceFieldUpdater, rather than creating a whole other object.